humalab 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

humalab/{scenario.py → scenarios/scenario.py}

@@ -1,6 +1,8 @@
 from typing import Any
+from threading import RLock
+
 import numpy as np
-from omegaconf import OmegaConf, ListConfig, AnyNode, DictConfig
+from omegaconf import OmegaConf, DictConfig, ListConfig
 import yaml
 from humalab.dists.bernoulli import Bernoulli
 from humalab.dists.categorical import Categorical
@@ -10,9 +12,7 @@ from humalab.dists.log_uniform import LogUniform
 from humalab.dists.gaussian import Gaussian
 from humalab.dists.truncated_gaussian import TruncatedGaussian
 from functools import partial
-from humalab.constants import EpisodeStatus
-from humalab.metrics.dist_metric import DistributionMetric
-from humalab.metrics.metric import MetricGranularity
+from humalab.constants import GraphType, MetricDimType
 import copy
 import uuid
 
@@ -37,39 +37,40 @@ DISTRIBUTION_MAP = {
 
     # 2D distributions
     "uniform_2d": Uniform,
-    "bernoulli_2d": Bernoulli,
-    "categorical_2d": Categorical,
-    "discrete_2d": Discrete,
-    "log_uniform_2d": LogUniform,
+    # "bernoulli_2d": Bernoulli,
+    # "categorical_2d": Categorical,
+    # "discrete_2d": Discrete,
+    # "log_uniform_2d": LogUniform,
     "gaussian_2d": Gaussian,
     "truncated_gaussian_2d": TruncatedGaussian,
 
     # 3D distributions
     "uniform_3d": Uniform,
-    "bernoulli_3d": Bernoulli,
-    "categorical_3d": Categorical,
-    "discrete_3d": Discrete,
-    "log_uniform_3d": LogUniform,
+    # "bernoulli_3d": Bernoulli,
+    # "categorical_3d": Categorical,
+    # "discrete_3d": Discrete,
+    # "log_uniform_3d": LogUniform,
     "gaussian_3d": Gaussian,
     "truncated_gaussian_3d": TruncatedGaussian,
 
     # 4D distributions
-    "uniform_4d": Uniform,
-    "bernoulli_4d": Bernoulli,
-    "categorical_4d": Categorical,
-    "discrete_4d": Discrete,
-    "log_uniform_4d": LogUniform,
-    "gaussian_4d": Gaussian,
-    "truncated_gaussian_4d": TruncatedGaussian,
-
+    # "uniform_4d": Uniform,
+    # "bernoulli_4d": Bernoulli,
+    # "categorical_4d": Categorical,
+    # "discrete_4d": Discrete,
+    # "log_uniform_4d": LogUniform,
+    # "gaussian_4d": Gaussian,
+    # "truncated_gaussian_4d": TruncatedGaussian,
+    
     # nD distributions
-    "uniform_nd": Uniform,
-    "bernoulli_nd": Bernoulli,
-    "categorical_nd": Categorical,
-    "discrete_nd": Discrete,
-    "log_uniform_nd": LogUniform,
-    "gaussian_nd": Gaussian,
-    "truncated_gaussian_nd": TruncatedGaussian,
+    # "uniform_nd": Uniform,
+    # "bernoulli_nd": Bernoulli,
+    # "categorical_nd": Categorical,
+    # "discrete_nd": Discrete,
+    # "log_uniform_nd": LogUniform,
+    # "gaussian_nd": Gaussian,
+    # "truncated_gaussian_nd": TruncatedGaussian,
+    
 }
 
 DISTRIBUTION_DIMENSION_MAP = {
@@ -93,39 +94,13 @@ DISTRIBUTION_DIMENSION_MAP = {
 
     # 2D distributions
     "uniform_2d": 2,
-    "bernoulli_2d": 2,
-    "categorical_2d": 2,
-    "discrete_2d": 2,
-    "log_uniform_2d": 2,
     "gaussian_2d": 2,
     "truncated_gaussian_2d": 2,
 
     # 3D distributions
     "uniform_3d": 3,
-    "bernoulli_3d": 3,
-    "categorical_3d": 3,
-    "discrete_3d": 3,
-    "log_uniform_3d": 3,
     "gaussian_3d": 3,
     "truncated_gaussian_3d": 3,
-
-    # 4D distributions
-    "uniform_4d": 4,
-    "bernoulli_4d": 4,
-    "categorical_4d": 4,
-    "discrete_4d": 4,
-    "log_uniform_4d": 4,
-    "gaussian_4d": 4,
-    "truncated_gaussian_4d": 4,
-
-    # nD distributions
-    "uniform_nd": -1,
-    "bernoulli_nd": -1,
-    "categorical_nd": -1,
-    "discrete_nd": -1,
-    "log_uniform_nd": -1,
-    "gaussian_nd": -1,
-    "truncated_gaussian_nd": -1,
 }
 
 DISTRIBUTION_PARAM_NUM_MAP = {
@@ -149,39 +124,73 @@ DISTRIBUTION_PARAM_NUM_MAP = {
 
     # 2D distributions
     "uniform_2d": 2,
-    "bernoulli_2d": 1,
-    "categorical_2d": 2,
-    "discrete_2d": 3,
-    "log_uniform_2d": 2,
     "gaussian_2d": 2,
     "truncated_gaussian_2d": 4,
 
     # 3D distributions
     "uniform_3d": 2,
-    "bernoulli_3d": 1,
-    "categorical_3d": 2,
-    "discrete_3d": 3,
-    "log_uniform_3d": 2,
     "gaussian_3d": 2,
     "truncated_gaussian_3d": 4,
+}
 
-    # 4D distributions
-    "uniform_4d": 2,
-    "bernoulli_4d": 1,
-    "categorical_4d": 2,
-    "discrete_4d": 3,
-    "log_uniform_4d": 2,
-    "gaussian_4d": 2,
-    "truncated_gaussian_4d": 4,
+SCENARIO_STATS_DIM_TYPE_MAP = {
+    # 0D distributions
+    "uniform": MetricDimType.ONE_D,
+    "bernoulli": MetricDimType.ONE_D,
+    "categorical": MetricDimType.ONE_D,
+    "discrete": MetricDimType.ONE_D,
+    "log_uniform": MetricDimType.ONE_D,
+    "gaussian": MetricDimType.ONE_D,
+    "truncated_gaussian": MetricDimType.ONE_D,
 
-    # nD distributions
-    "uniform_nd": 3,
-    "bernoulli_nd": 2,
-    "categorical_nd": 3,
-    "discrete_nd": 4,
-    "log_uniform_nd": 3,
-    "gaussian_nd": 3,
-    "truncated_gaussian_nd": 5,
+    # 1D distributions
+    "uniform_1d": MetricDimType.ONE_D,
+    "bernoulli_1d": MetricDimType.ONE_D,
+    "categorical_1d": MetricDimType.ONE_D,
+    "discrete_1d": MetricDimType.ONE_D,
+    "log_uniform_1d": MetricDimType.ONE_D,
+    "gaussian_1d": MetricDimType.ONE_D,
+    "truncated_gaussian_1d": MetricDimType.ONE_D,
+
+    # 2D distributions
+    "uniform_2d": MetricDimType.TWO_D,
+    "gaussian_2d": MetricDimType.TWO_D,
+    "truncated_gaussian_2d": MetricDimType.TWO_D,
+
+    # 3D distributions
+    "uniform_3d": MetricDimType.THREE_D,
+    "gaussian_3d": MetricDimType.THREE_D,
+    "truncated_gaussian_3d": MetricDimType.THREE_D,
+}
+
+DISTRIBUTION_GRAPH_TYPE = {
+    # 0D distributions
+    "uniform": GraphType.HISTOGRAM,
+    "bernoulli": GraphType.BAR,
+    "categorical": GraphType.BAR,
+    "discrete": GraphType.BAR,
+    "log_uniform": GraphType.HISTOGRAM,
+    "gaussian": GraphType.GAUSSIAN,
+    "truncated_gaussian": GraphType.GAUSSIAN,
+
+    # 1D distributions
+    "uniform_1d": GraphType.HISTOGRAM,
+    "bernoulli_1d": GraphType.BAR,
+    "categorical_1d": GraphType.BAR,
+    "discrete_1d": GraphType.BAR,
+    "log_uniform_1d": GraphType.HISTOGRAM,
+    "gaussian_1d": GraphType.GAUSSIAN,
+    "truncated_gaussian_1d": GraphType.GAUSSIAN,
+
+    # 2D distributions
+    "uniform_2d": GraphType.SCATTER,
+    "gaussian_2d": GraphType.HEATMAP,
+    "truncated_gaussian_2d": GraphType.HEATMAP,
+
+    # 3D distributions
+    "uniform_3d": GraphType.THREE_D_MAP,
+    "gaussian_3d": GraphType.THREE_D_MAP,
+    "truncated_gaussian_3d": GraphType.THREE_D_MAP,
 }
 
 class Scenario:
@@ -192,9 +201,10 @@ class Scenario:
         self._cur_scenario = OmegaConf.create()
         self._scenario_id = None
 
+        self._episode_vals = {}
+        self._lock = RLock()
+
     def init(self,
-             run_id: str,
-             episode_id: str,
              scenario: str | list | dict | None = None, 
              seed: int | None=None, 
              scenario_id: str | None=None,
@@ -204,24 +214,30 @@ class Scenario:
         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 = None # num_env
+
+        # Parse scenario id
+        scenario_version = 1
+        if scenario_id is not None:
+            scenario_arr = scenario_id.split(":")
+            if len(scenario_arr) < 1:
+                raise ValueError("Invalid scenario_id format. Expected 'scenario_id' or 'scenario_name:version'.")
+            scenario_id = scenario_arr[0]
+            scenario_version = int(scenario_arr[1]) if len(scenario_arr) > 1 else None
         self._scenario_id = scenario_id or str(uuid.uuid4())
+        self._scenario_version = scenario_version
+
         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 _validate_distribution_params(self, dist_name: str, *args: tuple) -> None:
         dimensions = DISTRIBUTION_DIMENSION_MAP[dist_name]
@@ -291,13 +307,6 @@ class Scenario:
             root_yaml = yaml.safe_load(OmegaConf.to_yaml(_root_))
             key_path = self._get_node_path(root_yaml, str(_node_))
             
-            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 DISTRIBUTION_DIMENSION_MAP[dist_name] == -1:
@@ -315,7 +324,13 @@ class Scenario:
 
             if isinstance(ret_val, list):
                 ret_val = ListConfig(ret_val)
-            self._metrics[key_path].log(ret_val)
+            
+            self._episode_vals[key_path] = {
+                    "value": ret_val,
+                    "distribution": dist_name,
+                    "graph_type": DISTRIBUTION_GRAPH_TYPE[dist_name],
+                    "metric_dim_type": SCENARIO_STATS_DIM_TYPE_MAP[dist_name],
+                }
             return ret_val
 
         for dist_name in DISTRIBUTION_MAP.keys():
@@ -325,33 +340,18 @@ class Scenario:
         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.
+    def resolve(self) -> tuple[DictConfig | ListConfig, dict]:
+        """Resolve the scenario configuration, sampling all distributions.
+
+        Returns:
+            tuple[DictConfig | ListConfig, dict]: The resolved scenario and episode values.
         """
-        for metric in self._metrics.values():
-            metric.finish()
+        with self._lock:
+            cur_scenario = copy.deepcopy(self._scenario_template)
+            self._episode_vals = {}
+            OmegaConf.resolve(cur_scenario)
+            episode_vals = copy.deepcopy(self._episode_vals)
+            return cur_scenario, episode_vals
 
     @property
     def template(self) -> Any:
@@ -362,15 +362,6 @@ class Scenario:
         """
         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.
@@ -378,4 +369,4 @@ class Scenario:
         Returns:
             str: The current scenario as a YAML string.
         """
-        return OmegaConf.to_yaml(self._cur_scenario)
+        return OmegaConf.to_yaml(self._scenario_template)

humalab/scenarios/scenario_operator.py

@@ -0,0 +1,82 @@
+from typing import Optional
+from dataclasses import dataclass
+
+from humalab.humalab_api_client import HumaLabApiClient
+from humalab.scenarios.scenario import Scenario
+from humalab.constants import DEFAULT_PROJECT
+
+@dataclass
+class ScenarioMetadata:
+    id: str
+    version: int
+    project: str
+    name: str
+    description: str | None
+    created_at: str
+    updated_at: str
+
+
+def list_scenarios(project: str = DEFAULT_PROJECT,
+                   limit: int = 20,
+                   offset: int = 0,     
+                   include_inactive: bool = False,
+                   search: Optional[str] = None,
+                   status_filter: Optional[str] = None,
+
+                   base_url: str | None = None,
+                   api_key: str | None = None,
+                   timeout: float | None = None,
+                   ) -> list[dict]:
+    """
+    List all scenarios for a given project.
+    
+    Args:
+        project: The project name to list scenarios from.
+        base_url: The base URL of the HumaLab API.
+        api_key: The API key for authentication.
+        timeout: The timeout for API requests.
+
+    Returns:
+        A list of scenario metadata dictionaries.
+    """
+    api_client = HumaLabApiClient(base_url=base_url,
+                                  api_key=api_key,
+                                  timeout=timeout)
+    resp = api_client.get_scenarios(project_name=project,
+                                    limit=limit,
+                                    offset=offset,
+                                    include_inactive=include_inactive,
+                                    search=search,
+                                    status_filter=status_filter)
+    ret_list = []
+    for scenario in resp.get("scenarios", []):
+        scenario["project"] = project
+        ret_list.append(ScenarioMetadata(id=scenario["uuid"],
+                                         version=scenario["version"],
+                                         project=project,
+                                         name=scenario["name"],
+                                         description=scenario.get("description"),
+                                         created_at=scenario.get("created_at"),
+                                         updated_at=scenario.get("updated_at")))
+    return ret_list
+
+def get_scenario(scenario_id: str,
+                 version: int | None = None,
+                 project: str = DEFAULT_PROJECT,
+                 seed: int | None=None,
+                 
+                 base_url: str | None = None,
+                 api_key: str | None = None,
+                 timeout: float | None = None,) -> Scenario:
+    api_client = HumaLabApiClient(base_url=base_url,
+                                  api_key=api_key,
+                                  timeout=timeout)
+    scenario_resp = api_client.get_scenario(
+        project_name=project,
+        uuid=scenario_id, version=version)
+    scenario = Scenario()
+
+    scenario.init(scenario=scenario_resp["yaml_content"],
+                  seed=seed,
+                  scenario_id=f"{scenario_id}:{version}" if version is not None else scenario_id)
+    return scenario