humalab 0.0.1__py3-none-any.whl

humalab/metrics/summary.py

@@ -0,0 +1,54 @@
+
+from humalab_sdk.metrics.metric import MetricGranularity, Metrics, MetricType
+from humalab_sdk.constants import EpisodeStatus
+
+
+class Summary(Metrics):
+    def __init__(self, 
+                 name: str, 
+                 summary: str,
+                 episode_id: str,
+                 run_id: str,
+                 granularity: MetricGranularity = MetricGranularity.RUN,
+                 ) -> None:
+        """
+        A summary metric that captures a single value per episode or run.
+
+        Args:
+            name (str): The name of the metric.
+            summary (str | None): Specify aggregate metrics added to summary.
+                Supported aggregations include "min", "max", "mean", "last",
+                "first", and "none". "none" prevents a summary
+                from being generated.
+            granularity (MetricGranularity): The granularity of the metric.
+        """
+        if granularity == MetricGranularity.RUN:
+            raise ValueError("Summary metrics cannot have RUN granularity.")
+        if summary not in {"min", "max", "mean", "last", "first", "none"}:
+            raise ValueError(f"Unsupported summary type: {summary}. Supported types are 'min', 'max', 'mean', 'last', 'first', and 'none'.")
+        super().__init__(name, MetricType.SUMMARY, episode_id=episode_id, run_id=run_id, granularity=granularity)
+        self.summary = summary
+
+    def _submit(self) -> None:
+        if not self._values:
+            return
+        # For summary metrics, we only keep the latest value
+        if self.summary == "last":
+            self._values = [self._values[-1]]
+        elif self.summary == "first":
+            self._values = [self._values[0]]
+        elif self.summary == "none":
+            self._values = []
+        elif self.summary in {"min", "max", "mean"}:
+            if not self._values:
+                self._values = []
+            else:
+                if self.summary == "min":
+                    agg_value = min(self._values)
+                elif self.summary == "max":
+                    agg_value = max(self._values)
+                elif self.summary == "mean":
+                    agg_value = sum(self._values) / len(self._values)
+                self._values = [agg_value]
+
+        super()._submit()

humalab/run.py

@@ -0,0 +1,214 @@
+import uuid
+from humalab_sdk.metrics.dist_metric import DistributionMetric
+from humalab_sdk.metrics.metric import MetricGranularity, MetricType, Metrics
+from humalab_sdk.constants import EpisodeStatus
+
+from humalab_sdk.metrics.summary import Summary
+from humalab_sdk.scenario import Scenario
+
+class Run:
+    def __init__(self,
+                 entity: str,
+                 project: str,
+                 scenario: Scenario,
+                 name: str | None = None,
+                 description: str | None = None,
+                 id: str | None = None,
+                 tags: list[str] | None = None,
+                 ) -> None:
+        """
+        Initialize a new Run instance.
+        
+        Args:
+            entity (str): The entity (user or team) under which the run is created.
+            project (str): The project name under which the run is created.
+            scenario (Scenario): The scenario instance for the run.
+            name (str | None): The name of the run.
+            description (str | None): A description of the run.
+            id (str | None): The unique identifier for the run. If None, a UUID is generated.
+            tags (list[str] | None): A list of tags associated with the run.
+        """
+        self._entity = entity
+        self._project = project
+        self._id = id or str(uuid.uuid4())
+        self._name = name or ""
+        self._description = description or ""
+        self._tags = tags or []
+        self._finished = False
+
+        self._episode = str(uuid.uuid4())
+
+        self._scenario = scenario
+
+        self._metrics = {}
+    
+    @property
+    def entity(self) -> str:
+        """The entity (user or team) under which the run is created.
+        
+        Returns:
+            str: The entity name.
+        """
+        return self._entity
+    
+    @property
+    def project(self) -> str:
+        """The project name under which the run is created.
+        
+        Returns:
+            str: The project name.
+        """
+        return self._project
+    
+    @property
+    def id(self) -> str:
+        """The unique identifier for the run.
+        
+        Returns:
+            str: The run ID.
+        """
+        return self._id
+    
+    @property
+    def name(self) -> str:
+        """The name of the run.
+
+        Returns:
+            str: The run name.
+        """
+        return self._name
+    
+    @property
+    def description(self) -> str:
+        """The description of the run.
+
+        Returns:
+            str: The run description.
+        """
+        return self._description
+    
+    @property
+    def tags(self) -> list[str]:
+        """The tags associated with the run.
+
+        Returns:
+            list[str]: The list of tags.
+        """
+        return self._tags
+    
+    @property
+    def episode(self) -> str:
+        """The episode ID for the run.
+
+        Returns:
+            str: The episode ID.
+        """
+        return self._episode
+    
+    @property
+    def scenario(self) -> Scenario:
+        """The scenario associated with the run.
+
+        Returns:
+            Scenario: The scenario instance.
+        """
+        return self._scenario
+
+    def finish(self,
+               status: EpisodeStatus = EpisodeStatus.PASS,
+               quiet: bool | None = None) -> None:
+        """Finish the run and submit final metrics.
+
+        Args:
+            status (EpisodeStatus): The final status of the episode.
+            quiet (bool | None): Whether to suppress output.
+        """
+        self._finished = True
+        self._scenario.finish()
+        for metric in self._metrics.values():
+            metric.finish(status=status)
+    
+    def log(self,
+            data: dict,
+            step: int | None = None,
+            commit: bool = True,
+            ) -> None:
+        """Log metrics for the run.
+
+        Args:
+            data (dict): A dictionary of metric names and their values.
+            step (int | None): The step number for the metrics.
+            commit (bool): Whether to commit the metrics immediately.
+        """
+        for key, value in data.items():
+            if key in self._metrics:
+                metric = self._metrics[key]
+                metric.log(value, step=step, commit=commit)
+            else:
+                self._metrics[key] = Metrics(key, 
+                                             metric_type=MetricType.DEFAULT,
+                                             run_id=self._id,
+                                             granularity=MetricGranularity.EPISODE,
+                                             episode_id=self._episode)
+                self._metrics[key].log(value, step=step, commit=commit)
+
+    def reset(self, status: EpisodeStatus = EpisodeStatus.PASS) -> None:
+        """Reset the run for a new episode.
+
+        Args:
+            status (EpisodeStatus): The status of the current episode before reset.
+        """
+        self._submit_episode_status(status=status, episode=self._episode)
+        self._episode = str(uuid.uuid4())
+        self._finished = False
+        self._scenario.reset(episode_id=self._episode)
+        for metric in self._metrics.values():
+            metric.reset(episode=self._episode)
+    
+    def _submit_episode_status(self, status: EpisodeStatus, episode: str) -> None:
+        # TODO: Implement submission of episode status
+        pass
+
+    def define_metric(self, 
+                      name: str, 
+                      metric_type: MetricType = MetricType.DEFAULT,
+                      granularity: MetricGranularity = MetricGranularity.RUN,
+                      distribution_type: str | None = None,
+                      summary: str | None = None,
+                      replace: bool = False) -> None:
+        """Define a new metric for the run.
+        
+        Args:
+            name (str): The name of the metric.
+            metric_type (MetricType): The type of the metric.
+            granularity (MetricGranularity): The granularity of the metric.
+            distribution_type (str | None): The type of distribution if metric_type is DISTRIBUTION.
+            summary (str | None): Specify aggregate metrics added to summary.
+                Supported aggregations include "min", "max", "mean", "last",
+                "first", and "none". "none" prevents a summary
+                from being generated.
+            replace (bool): Whether to replace the metric if it already exists.
+        """
+        if name not in self._metrics or replace:
+            if metric_type == MetricType.DISTRIBUTION:
+                if distribution_type is None:
+                    raise ValueError("distribution_type must be specified for distribution metrics.")
+                self._metrics[name] = DistributionMetric(name=name, 
+                                                         distribution_type=distribution_type, 
+                                                         run_id=self._id,
+                                                         episode_id=self._episode,
+                                                         granularity=granularity)  
+            elif summary is not None:
+                self._metrics[name] = Summary(name=name, 
+                                              summary=summary, 
+                                              run_id=self._id,
+                                              episode_id=self._episode,
+                                              granularity=granularity)
+            else:
+                self._metrics[name] = Metrics(name=name, 
+                                              metric_type=metric_type, 
+                                              run_id=self._id,
+                                              episode_id=self._episode,
+                                              granularity=granularity)
+        else:
+            raise ValueError(f"Metric {name} already exists.")

humalab/scenario.py

@@ -0,0 +1,225 @@
+from typing import Any
+import numpy as np
+from omegaconf import OmegaConf, ListConfig, AnyNode, DictConfig
+import yaml
+from humalab_sdk.dists.bernoulli import Bernoulli
+from humalab_sdk.dists.categorical import Categorical
+from humalab_sdk.dists.uniform import Uniform
+from humalab_sdk.dists.discrete import Discrete
+from humalab_sdk.dists.log_uniform import LogUniform
+from humalab_sdk.dists.gaussian import Gaussian
+from humalab_sdk.dists.truncated_gaussian import TruncatedGaussian
+from functools import partial
+from humalab_sdk.constants import EpisodeStatus
+from humalab_sdk.metrics.dist_metric import DistributionMetric
+from humalab_sdk.metrics.metric import MetricGranularity
+import copy
+import uuid
+
+DISTRIBUTION_MAP = {
+    "uniform": Uniform,
+    "bernoulli": Bernoulli,
+    "categorical": Categorical,
+    "discrete": Discrete,
+    "log_uniform": LogUniform,
+    "gaussian": Gaussian,
+    "truncated_gaussian": TruncatedGaussian,
+}
+
+DISTRIBUTION_PARAM_NUM_MAP = {
+    "uniform": 3,
+    "bernoulli": 2,
+    "categorical": 3,
+    "discrete": 4,
+    "log_uniform": 3,
+    "gaussian": 3,
+    "truncated_gaussian": 5,
+}
+
+class Scenario:
+    dist_cache = {}
+    def __init__(self) -> None:
+        self._generator = np.random.default_rng()
+        self._scenario_template = OmegaConf.create()
+        self._cur_scenario = OmegaConf.create()
+        self._scenario_id = None
+
+    def init(self,
+             run_id: str,
+             episode_id: str,
+             scenario: str | list | dict | None = None, 
+             seed: int | None=None, 
+             scenario_id: str | None=None,
+             num_env: int | None = None) -> None:
+        """
+        Initialize the scenario with the given parameters.
+        
+        Args:
+            run_id: The ID of the current run.
+            episode_id: The ID of the current episode.
+            scenario: The scenario configuration (YAML string, list, or dict).
+            seed: Optional seed for random number generation.
+            scenario_id: Optional scenario ID. If None, a new UUID is generated.
+            num_env: Optional number of parallel environments.
+        """
+        self._run_id = run_id
+        self._episode_id = episode_id
+        self._metrics = {}
+
+        self._num_env = num_env
+        self._scenario_id = scenario_id or str(uuid.uuid4())
+        self._generator = np.random.default_rng(seed)
+        self._configure()
+        scenario = scenario or {}
+        self._scenario_template = OmegaConf.create(scenario)
+        self.reset(episode_id=episode_id)
+
+    def _get_final_size(self, size: int | tuple[int, ...] | None) -> int | tuple[int, ...] | None:
+        n = self._num_env
+        if size is None:
+            return n
+        if n is None:
+            return size
+        if isinstance(size, int):
+            return (n, size)
+        return (n, *size)
+    
+    def _get_node_path(self, root: dict | list, node: str) -> str:
+        if isinstance(root, list):
+            root = {str(i): v for i, v in enumerate(root)}
+        
+        for key, value in root.items():
+            if value == node:
+                return str(key)
+            if isinstance(value, dict):
+                sub_path = self._get_node_path(value, node)
+                if sub_path:
+                    return f"{key}.{sub_path}"
+            elif isinstance(value, list):
+                for idx, item in enumerate(value):
+                    if item == node:
+                        return f"{key}[{idx}]"
+                    if isinstance(item, (dict, list)):
+                        sub_path = self._get_node_path(item, node)
+                        if sub_path:
+                            return f"{key}[{idx}].{sub_path}"
+        return ""
+
+    @staticmethod
+    def _convert_to_python(obj) -> Any:
+        if not isinstance(obj, (np.ndarray, np.generic)):
+            return obj
+
+        # NumPy scalar (np.generic) or 0-D ndarray
+        if isinstance(obj, np.generic) or (isinstance(obj, np.ndarray) and obj.ndim == 0):
+            return obj.item()
+
+        # Regular ndarray (1-D or higher)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+
+        return obj
+
+    def _configure(self) -> None:
+        self._clear_resolvers()
+        def distribution_resolver(dist_name: str, *args, _node_, _root_, _parent_, **kwargs):
+            if len(args) > DISTRIBUTION_PARAM_NUM_MAP[dist_name]:
+                args = args[:DISTRIBUTION_PARAM_NUM_MAP[dist_name]]
+                print(f"Warning: Too many parameters for {dist_name}, expected {DISTRIBUTION_PARAM_NUM_MAP[dist_name]}, got {len(args)}. Extra parameters will be ignored.")
+            
+            # print("_node_: ", _node_, type(_node_))
+            # print("_root_: ", _root_, type(_root_))
+            # print("_parent_: ", _parent_, type(_parent_))
+            # print("Args: ", args, len(args))
+            # print("Kwargs: ", kwargs)
+
+            root_yaml = yaml.safe_load(OmegaConf.to_yaml(_root_))
+            key_path = self._get_node_path(root_yaml, str(_node_))
+            # print("Key path: ", key_path)
+            
+            if key_path not in self._metrics:
+                self._metrics[key_path] = DistributionMetric(name=key_path,
+                                                            distribution_type=dist_name,
+                                                            run_id=self._run_id,
+                                                            episode_id=self._episode_id,
+                                                            granularity=MetricGranularity.EPISODE)
+
+            shape = None 
+            
+            if len(args) == DISTRIBUTION_PARAM_NUM_MAP[dist_name]:
+                shape = args[DISTRIBUTION_PARAM_NUM_MAP[dist_name] - 1]
+                args = args[:-1]
+            shape = self._get_final_size(shape)
+
+            key = str(_node_)
+            if key not in Scenario.dist_cache:
+                Scenario.dist_cache[key] = DISTRIBUTION_MAP[dist_name].create(self._generator, *args, size=shape, **kwargs)
+            ret_val = Scenario.dist_cache[key].sample()
+            ret_val = Scenario._convert_to_python(ret_val)
+
+            if isinstance(ret_val, list):
+                ret_val = ListConfig(ret_val)
+            self._metrics[key_path].log(ret_val)
+            return ret_val
+
+        for dist_name in DISTRIBUTION_MAP.keys():
+            OmegaConf.register_new_resolver(dist_name, partial(distribution_resolver, dist_name))
+
+    def _clear_resolvers(self) -> None:
+        self.dist_cache.clear()
+        OmegaConf.clear_resolvers()
+    
+    def __getattr__(self, name: Any) -> Any:
+        if name in self._cur_scenario:
+            return self._cur_scenario[name]
+        raise AttributeError(f"'Scenario' object has no attribute '{name}'")
+
+    def __getitem__(self, key: Any) -> Any:
+        if key in self._cur_scenario:
+            return self._cur_scenario[key]
+        raise KeyError(f"'Scenario' object has no key '{key}'")
+
+    def reset(self,
+              episode_id: str | None = None) -> None:
+        """Reset the scenario for a new episode.
+        
+        Args:
+            episode_id: Optional new episode ID. If None, keeps the current episode ID.
+        """
+        for metric in self._metrics.values():
+            metric.reset(episode_id=episode_id)
+        self._cur_scenario = copy.deepcopy(self._scenario_template)
+        OmegaConf.resolve(self._cur_scenario)
+    
+    def finish(self) -> None:
+        """Finish the scenario and submit final metrics.
+        """
+        for metric in self._metrics.values():
+            metric.finish()
+
+    @property
+    def template(self) -> Any:
+        """The template scenario configuration.
+        
+        Returns:
+            Any: The template scenario as an OmegaConf object.
+        """
+        return self._scenario_template
+    
+    @property
+    def cur_scenario(self) -> Any:
+        """The current scenario configuration.
+
+        Returns:
+            Any: The current scenario as an OmegaConf object.
+        """
+        return self._cur_scenario
+    
+    @property
+    def yaml(self) -> str:
+        """The current scenario configuration as a YAML string.
+
+        Returns:
+            str: The current scenario as a YAML string.
+        """
+        return OmegaConf.to_yaml(self._cur_scenario)