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

humalab/dists/gaussian.py

@@ -4,6 +4,12 @@ import numpy as np
 
 
 class Gaussian(Distribution):
+    """Gaussian (normal) distribution.
+
+    Samples values from a normal distribution with specified mean (loc) and
+    standard deviation (scale). Supports scalar outputs as well as multi-dimensional
+    arrays with 1D, 2D, or 3D variants.
+    """
     def __init__(self,
                  generator: np.random.Generator,
                  loc: float | Any,
@@ -25,6 +31,15 @@ class Gaussian(Distribution):
 
     @staticmethod
     def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (loc, scale).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
         arg1 = args[0]
         arg2 = args[1]
         if dimensions == 0:
@@ -37,18 +52,28 @@ class Gaussian(Distribution):
             return True
         if not isinstance(arg1, (int, float)):
             if isinstance(arg1, (list, np.ndarray)):
-                if len(arg1) > dimensions:
+                if len(arg1) != dimensions:
                     return False
         if not isinstance(arg2, (int, float)):
             if isinstance(arg2, (list, np.ndarray)):
-                if len(arg2) > dimensions:
+                if len(arg2) != dimensions:
                     return False
         return True
 
     def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the Gaussian distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) from N(loc, scale).
+        """
         return self._generator.normal(loc=self._loc, scale=self._scale, size=self._size)
 
     def __repr__(self) -> str:
+        """String representation of the Gaussian distribution.
+
+        Returns:
+            str: String representation showing loc, scale, and size.
+        """
         return f"Gaussian(loc={self._loc}, scale={self._scale}, size={self._size})"
     
     @staticmethod

humalab/dists/log_uniform.py

@@ -4,6 +4,13 @@ from typing import Any
 import numpy as np
 
 class LogUniform(Distribution):
+    """Log-uniform distribution.
+
+    Samples values uniformly in log-space, useful for hyperparameters that
+    span multiple orders of magnitude (e.g., learning rates). The result is
+    exp(uniform(log(low), log(high))). Supports scalar outputs as well as
+    multi-dimensional arrays with 1D variants.
+    """
     def __init__(self, 
                  generator: np.random.Generator,
                  low: float | Any,
@@ -25,6 +32,15 @@ class LogUniform(Distribution):
     
     @staticmethod
     def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
         arg1 = args[0]
         arg2 = args[1]
         if dimensions == 0:
@@ -37,18 +53,28 @@ class LogUniform(Distribution):
             return True
         if not isinstance(arg1, (int, float)):
             if isinstance(arg1, (list, np.ndarray)):
-                if len(arg1) > dimensions:
+                if len(arg1) != dimensions:
                     return False
         if not isinstance(arg2, (int, float)):
             if isinstance(arg2, (list, np.ndarray)):
-                if len(arg2) > dimensions:
+                if len(arg2) != dimensions:
                     return False
         return True
 
     def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the log-uniform distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) in log-space.
+        """
         return np.exp(self._generator.uniform(self._log_low, self._log_high, size=self._size))
-    
+
     def __repr__(self) -> str:
+        """String representation of the log-uniform distribution.
+
+        Returns:
+            str: String representation showing low, high, and size.
+        """
         return f"LogUniform(low={np.exp(self._log_low)}, high={np.exp(self._log_high)}, size={self._size})"
 
     @staticmethod

humalab/dists/truncated_gaussian.py

@@ -4,6 +4,13 @@ import numpy as np
 
 
 class TruncatedGaussian(Distribution):
+    """Truncated Gaussian (normal) distribution.
+
+    Samples values from a normal distribution with specified mean (loc) and
+    standard deviation (scale), but constrained to lie within [low, high].
+    Values outside the bounds are resampled until they fall within range.
+    Supports scalar outputs as well as multi-dimensional arrays with 1D, 2D, or 3D variants.
+    """
     def __init__(self,
                  generator: np.random.Generator,
                  loc: float | Any,
@@ -31,6 +38,15 @@ class TruncatedGaussian(Distribution):
 
     @staticmethod
     def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (loc, scale, low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
         arg1 = args[0]
         arg2 = args[1]
         arg3 = args[2]
@@ -49,23 +65,31 @@ class TruncatedGaussian(Distribution):
             return True
         if not isinstance(arg1, (int, float)):
             if isinstance(arg1, (list, np.ndarray)):
-                if len(arg1) > dimensions:
+                if len(arg1) != dimensions:
                     return False
         if not isinstance(arg2, (int, float)):
             if isinstance(arg2, (list, np.ndarray)):
-                if len(arg2) > dimensions:
+                if len(arg2) != dimensions:
                     return False
         if not isinstance(arg3, (int, float)):
             if isinstance(arg3, (list, np.ndarray)):
-                if len(arg3) > dimensions:
+                if len(arg3) != dimensions:
                     return False
         if not isinstance(arg4, (int, float)):
             if isinstance(arg4, (list, np.ndarray)):
-                if len(arg4) > dimensions:
+                if len(arg4) != dimensions:
                     return False
         return True
     
     def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the truncated Gaussian distribution.
+
+        Samples are generated from N(loc, scale) and resampled if they fall
+        outside [low, high].
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) within [low, high].
+        """
         samples = self._generator.normal(loc=self._loc, scale=self._scale, size=self._size)
         mask = (samples < self._low) | (samples > self._high)
         while np.any(mask):
@@ -74,6 +98,11 @@ class TruncatedGaussian(Distribution):
         return samples
 
     def __repr__(self) -> str:
+        """String representation of the truncated Gaussian distribution.
+
+        Returns:
+            str: String representation showing loc, scale, low, high, and size.
+        """
         return f"TruncatedGaussian(loc={self._loc}, scale={self._scale}, low={self._low}, high={self._high}, size={self._size})"
     
     @staticmethod

humalab/dists/uniform.py

@@ -4,6 +4,11 @@ from typing import Any
 import numpy as np
 
 class Uniform(Distribution):
+    """Uniform distribution over a continuous or discrete range.
+
+    Samples values uniformly from the half-open interval [low, high). Supports
+    scalar outputs as well as multi-dimensional arrays with 1D, 2D, or 3D variants.
+    """
     def __init__(self, 
                  generator: np.random.Generator,
                  low: float | Any, 
@@ -25,6 +30,15 @@ class Uniform(Distribution):
 
     @staticmethod
     def validate(dimensions: int, *args) -> bool:
+        """Validate distribution parameters for the given dimensions.
+
+        Args:
+            dimensions (int): The number of dimensions (0 for scalar, -1 for any).
+            *args: The distribution parameters (low, high).
+
+        Returns:
+            bool: True if parameters are valid, False otherwise.
+        """
         arg1 = args[0]
         arg2 = args[1]
         if dimensions == 0:
@@ -46,9 +60,19 @@ class Uniform(Distribution):
         return True
 
     def _sample(self) -> int | float | np.ndarray:
+        """Generate a sample from the uniform distribution.
+
+        Returns:
+            int | float | np.ndarray: Sampled value(s) from [low, high).
+        """
         return self._generator.uniform(self._low, self._high, size=self._size)
 
     def __repr__(self) -> str:
+        """String representation of the uniform distribution.
+
+        Returns:
+            str: String representation showing low, high, and size.
+        """
         return f"Uniform(low={self._low}, high={self._high}, size={self._size})"
     
     @staticmethod

humalab/episode.py

@@ -1,20 +1,223 @@
+from humalab.constants import RESERVED_NAMES, ArtifactType
+from humalab.humalab_api_client import HumaLabApiClient, EpisodeStatus
+from humalab.metrics.code import Code
+from humalab.metrics.summary import Summary
+from humalab.metrics.metric import Metrics
+from omegaconf import DictConfig, ListConfig, OmegaConf
+from typing import Any
+import pickle
+import traceback
 
-from humalab.scenario import Scenario
-from omegaconf import DictConfig, OmegaConf
+from humalab.utils import is_standard_type
 
 
 class Episode:
-    def __init__(self, run_id: str, episode_id: str, scenario_conf: DictConfig):
-        self.run_id = run_id
-        self.episode_id = episode_id
-        self.scenario_conf = scenario_conf
+    """Represents a single episode within a run.
+
+    An Episode is a context manager that tracks a single execution instance of a
+    scenario. It provides access to scenario configuration values, supports metric
+    logging, and manages episode lifecycle with various completion statuses.
+
+    Episodes can be finished with different statuses:
+    - SUCCESS: Episode completed successfully
+    - FAILED: Episode failed
+    - CANCELED: Episode was discarded/canceled
+    - ERRORED: Episode encountered an error
+
+    Use as a context manager to automatically handle episode lifecycle:
+        with episode:
+            # Your code here
+            pass
+
+    Attributes:
+        run_id (str): The unique identifier of the parent run.
+        episode_id (str): The unique identifier for this episode.
+        scenario (DictConfig | ListConfig): The resolved scenario configuration.
+        status (EpisodeStatus): The current status of the episode.
+        episode_vals (dict): The sampled values from scenario distributions.
+        is_finished (bool): Whether the episode has been finalized.
+    """
+    def __init__(self, 
+                 run_id: str, 
+                 episode_id: str, 
+                 scenario_conf: DictConfig | ListConfig,
+                 episode_vals: dict | None = None,
+
+                 base_url: str | None = None,
+                 api_key: str | None = None,
+                 timeout: float | None = None,):
+        self._run_id = run_id
+        self._episode_id = episode_id
+        self._episode_status = EpisodeStatus.RUNNING
+        self._scenario_conf = scenario_conf
+        self._logs = {}
+        self._episode_vals = episode_vals or {}
+        self._is_finished = False
+
+        self._api_client = HumaLabApiClient(base_url=base_url,
+                                            api_key=api_key,
+                                            timeout=timeout)
 
     @property
-    def scenario(self) -> DictConfig:
-        return self.scenario_conf
+    def run_id(self) -> str:
+        """The unique identifier of the parent run.
 
-    def finish(self):
-        print(f"Finishing episode {self.episode_id} for scenario {self.scenario.name}")
+        Returns:
+            str: The run ID.
+        """
+        return self._run_id
+
+    @property
+    def episode_id(self) -> str:
+        """The unique identifier for this episode.
+
+        Returns:
+            str: The episode ID.
+        """
+        return self._episode_id
+
+    @property
+    def scenario(self) -> DictConfig | ListConfig:
+        """The resolved scenario configuration for this episode.
+
+        Returns:
+            DictConfig | ListConfig: The scenario configuration.
+        """
+        return self._scenario_conf
+
+    @property
+    def status(self) -> EpisodeStatus:
+        """The current status of the episode.
+
+        Returns:
+            EpisodeStatus: The episode status.
+        """
+        return self._episode_status
+
+    @property
+    def episode_vals(self) -> dict:
+        """The sampled values from scenario distributions.
+
+        Returns:
+            dict: Dictionary mapping scenario variable names to their sampled values.
+        """
+        return self._episode_vals
+
+    @property
+    def is_finished(self) -> bool:
+        """Whether the episode has been finalized.
+
+        Returns:
+            bool: True if the episode is finished, False otherwise.
+        """
+        return self._is_finished
+    
+    def __enter__(self):
+        """Enter the episode context."""
+        return self
+
+    def __exit__(self, exception_type, exception_value, exception_traceback):
+        """Exit the episode context and finalize the episode."""
+        if self._is_finished:
+            return
+        if exception_type is not None:
+            err_msg = "".join(traceback.format_exception(exception_type, exception_value, exception_traceback))
+            self.finish(status=EpisodeStatus.ERRORED, err_msg=err_msg)
+        else:
+            self.finish(status=EpisodeStatus.SUCCESS)
+
+    def __getattr__(self, name: Any) -> Any:
+        """Access scenario configuration values as attributes.
+
+        Allows accessing scenario configuration using dot notation (e.g., episode.my_param).
+
+        Args:
+            name (Any): The attribute/key name from scenario configuration.
+
+        Returns:
+            Any: The value from scenario configuration.
+
+        Raises:
+            AttributeError: If the attribute is not in scenario configuration.
+        """
+        if name in self._scenario_conf:
+            return self._scenario_conf[name]
+        raise AttributeError(f"'Scenario' object has no attribute '{name}'")
+
+    def __getitem__(self, key: Any) -> Any:
+        """Access scenario configuration values using subscript notation.
+
+        Allows accessing scenario configuration using bracket notation (e.g., episode['my_param']).
+
+        Args:
+            key (Any): The key name from scenario configuration.
+
+        Returns:
+            Any: The value from scenario configuration.
+
+        Raises:
+            KeyError: If the key is not in scenario configuration.
+        """
+        if key in self._scenario_conf:
+            return self._scenario_conf[key]
+        raise KeyError(f"'Scenario' object has no key '{key}'")
+
+    def add_metric(self, name: str, metric: Metrics) -> None:
+        """Add a metric to track for this episode.
+
+        Args:
+            name (str): The name of the metric.
+            metric (Metrics): The metric instance to add.
+
+        Raises:
+            ValueError: If the name is already used.
+        """
+        if name in self._logs:
+            raise ValueError(f"{name} is a reserved name and is not allowed.")
+        self._logs[name] = metric
+    
+    def log_code(self, key: str, code_content: str) -> None:
+        """Log code content as an artifact.
+
+        Args:
+            key (str): The key for the code artifact.
+            code_content (str): The code content to log.
+        """
+        if key in RESERVED_NAMES:
+            raise ValueError(f"{key} is a reserved name and is not allowed.")
+        self._logs[key] = Code(
+            run_id=self._run_id,
+            key=key,
+            code_content=code_content,
+            episode_id=self._episode_id
+        )
+
+    def log(self, data: dict, x: dict | None = None, replace: bool = False) -> None:
+        """Log data points or values for the episode.
+
+        Args:
+            data (dict): Dictionary of key-value pairs to log.
+            x (dict | None): Optional dictionary of x-axis values for each key.
+            replace (bool): Whether to replace existing values. Defaults to False.
+
+        Raises:
+            ValueError: If a key is reserved or logging fails.
+        """
+        for key, value in data.items():
+            if key in RESERVED_NAMES:
+                raise ValueError(f"{key} is a reserved name and is not allowed.")
+            if key not in self._logs:
+                self._logs[key] = value
+            else:
+                cur_val = self._logs[key]
+                if isinstance(cur_val, Metrics):
+                    cur_x = x.get(key) if x is not None else None
+                    cur_val.log(value, x=cur_x, replace=replace)
+                else:
+                    if replace:
+                        self._logs[key] = value
+                    else:
+                        raise ValueError(f"Cannot log value for key '{key}' as there is already a value logged.")
 
     @property
     def yaml(self) -> str:
@@ -23,4 +226,81 @@ class Episode:
         Returns:
             str: The current scenario as a YAML string.
         """
-        return OmegaConf.to_yaml(self.scenario_conf)
+        return OmegaConf.to_yaml(self._scenario_conf)
+    
+    def discard(self) -> None:
+        """Mark the episode as discarded/canceled."""
+        self._finish(EpisodeStatus.CANCELED)
+
+    def success(self) -> None:
+        """Mark the episode as successfully completed."""
+        self._finish(EpisodeStatus.SUCCESS)
+
+    def fail(self) -> None:
+        """Mark the episode as failed."""
+        self._finish(EpisodeStatus.FAILED)
+
+    def finish(self, status: EpisodeStatus, err_msg: str | None = None) -> None:
+        """Finish the episode with a specific status.
+
+        Args:
+            status (EpisodeStatus): The final status of the episode.
+            err_msg (str | None): Optional error message if the episode errored.
+        """
+        if self._is_finished:
+            return
+        self._is_finished = True
+        self._episode_status = status
+
+        self._api_client.upload_code(
+            artifact_key="scenario",
+            run_id=self._run_id,
+            episode_id=self._episode_id,
+            code_content=self.yaml
+        )
+
+        # TODO: submit final metrics
+        for key, value in self._logs.items():
+            if isinstance(value, Summary):
+                metric_val = value.finalize()
+                pickled = pickle.dumps(metric_val["value"])
+                self._api_client.upload_python(
+                    artifact_key=key,
+                    run_id=self._run_id,
+                    episode_id=self._episode_id,
+                    pickled_bytes=pickled
+                )
+            elif isinstance(value, Metrics):
+                metric_val = value.finalize()
+                pickled = pickle.dumps(metric_val)
+                self._api_client.upload_metrics(
+                    artifact_key=key,
+                    run_id=self._run_id,
+                    episode_id=self._episode_id,
+                    pickled_bytes=pickled,
+                    graph_type=value.graph_type.value,
+                )
+            elif isinstance(value, Code):
+                self._api_client.upload_code(
+                    artifact_key=value.key,
+                    run_id=value.run_id,
+                    episode_id=value.episode_id,
+                    code_content=value.code_content
+                )
+            else:
+                if not is_standard_type(value):
+                    raise ValueError(f"Value for key '{key}' is not a standard type.")
+                pickled = pickle.dumps(value)
+                self._api_client.upload_python(
+                    artifact_key=key,
+                    run_id=self._run_id,
+                    episode_id=self._episode_id,
+                    pickled_bytes=pickled
+                )
+        
+        self._api_client.update_episode(
+            run_id=self._run_id,
+            episode_id=self._episode_id,
+            status=status,
+            err_msg=err_msg
+        )