libinephany 0.13.1__py3-none-any.whl

libinephany/aws/s3_functions.py

@@ -0,0 +1,57 @@
+# ======================================================================================================================
+#
+# IMPORTS
+#
+# ======================================================================================================================
+
+import boto3
+from loguru import logger
+
+# ======================================================================================================================
+#
+# CONSTANTS
+#
+# ======================================================================================================================
+
+S3_URI_SCHEME = "s3://"
+S3_BUCKET_SEPARATOR = "/"
+S3_BOTO_CLIENT = "s3"
+
+# ======================================================================================================================
+#
+# FUNCTIONS
+#
+# ======================================================================================================================
+
+
+def parse_s3_url(s3_url: str) -> tuple[str, str]:
+    """
+    :param s3_url: S3 URL to parse and extract the bucket and blob names from.
+    :return: Tuple of:
+        - Bucket name.
+        - Blob name.
+    """
+
+    if s3_url.startswith(S3_URI_SCHEME):
+        s3_url = s3_url.replace(S3_URI_SCHEME, "")
+
+    parts = s3_url.split(S3_BUCKET_SEPARATOR, 1)
+    bucket = parts[0]
+    key = parts[1] if len(parts) > 1 else ""
+
+    return bucket, key
+
+
+def download_s3_file(s3_url: str, local_path: str) -> None:
+    """
+    :param s3_url: S3 URL to download data from.
+    :param local_path: Path to save the contents of the download to.
+    """
+
+    bucket, key = parse_s3_url(s3_url)
+
+    s3 = boto3.client(S3_BOTO_CLIENT)
+
+    logger.info(f"Downloading {s3_url} to {local_path}...")
+    s3.download_file(bucket, key, local_path)
+    logger.success("Done.")

libinephany/observations/observation_utils.py

@@ -0,0 +1,243 @@
+# ======================================================================================================================
+#
+# IMPORTS
+#
+# ======================================================================================================================
+
+from enum import Enum
+from itertools import chain
+from typing import Any, Callable
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.optim as optim
+
+from libinephany.pydantic_models.schemas.tensor_statistics import TensorStatistics
+from libinephany.utils import optim_utils
+
+# ======================================================================================================================
+#
+# CONSTANTS
+#
+# ======================================================================================================================
+
+EXP_AVERAGE = "exp_avg"
+
+# ======================================================================================================================
+#
+# CLASSES
+#
+# ======================================================================================================================
+
+
+class StatisticsCallStage(Enum):
+
+    ON_BATCH_END = "on_batch_end"
+    ON_OPTIMIZER_STEP = "on_optimizer_step"
+    ON_TRAIN_END = "on_train_end"
+
+    FORWARD_HOOK = "forward_hook"
+
+
+class StatisticStorageTypes(Enum):
+
+    TENSOR_STATISTICS = TensorStatistics.__name__
+    FLOAT = float.__name__
+    VECTOR = "vector"
+
+
+# ======================================================================================================================
+#
+# FUNCTIONS
+#
+# ======================================================================================================================
+
+
+def get_exponential_weighted_average(values: list[float]) -> float:
+    """
+    :param values: List of values to average via EWA.
+    :return: EWA of the given values.
+    """
+
+    exp_weighted_average = pd.Series(values).ewm(alpha=0.1).mean().iloc[-1]
+    assert isinstance(exp_weighted_average, float)
+
+    return exp_weighted_average
+
+
+def apply_averaging_function_to_tensor_statistics(
+    tensor_statistics: list[TensorStatistics], averaging_function: Callable[[list[float]], float]
+) -> TensorStatistics:
+    """
+    :param tensor_statistics: List of statistics models to average over.
+    :param averaging_function: Function to average the values with.
+    :return: TensorStatistics containing the average over all given tensor statistics.
+    """
+
+    fields = TensorStatistics.model_fields.keys()
+    averaged_metrics = {
+        field: averaging_function([getattr(statistics, field) for statistics in tensor_statistics]) for field in fields
+    }
+
+    return TensorStatistics(**averaged_metrics)
+
+
+def apply_averaging_function_to_dictionary_of_tensor_statistics(
+    data: dict[str, list[TensorStatistics]], averaging_function: Callable[[list[float]], float]
+) -> dict[str, TensorStatistics]:
+    """
+    :param data: Dictionary mapping parameter group names to list of TensorStatistics from that parameter group.
+    :param averaging_function: Function to average the values with.
+    :return: Dictionary mapping parameter group names to TensorStatistics averaged over all statistics in the given
+    TensorStatistics models.
+    """
+
+    return {
+        group: apply_averaging_function_to_tensor_statistics(
+            tensor_statistics=metrics, averaging_function=averaging_function
+        )
+        for group, metrics in data.items()
+    }
+
+
+def apply_averaging_function_to_dictionary_of_metric_lists(
+    data: dict[str, list[float]], averaging_function: Callable[[list[float]], float]
+) -> dict[str, float]:
+    """
+    :param data: Dictionary mapping parameter group names to list of metrics from that parameter group.
+    :param averaging_function: Function to average the values with.
+    :return: Dictionary mapping parameter group names to averages over all metrics from each parameter group.
+    """
+
+    return {group: averaging_function(metrics) for group, metrics in data.items()}
+
+
+def average_tensor_statistics(tensor_statistics: list[TensorStatistics]) -> TensorStatistics:
+    """
+    :param tensor_statistics: List of TensorStatistics models to average into one model.
+    :return: Averages over all given tensor statistics models.
+    """
+
+    averaged = {
+        field: sum([getattr(statistics_model, field) for statistics_model in tensor_statistics])
+        for field in TensorStatistics.model_fields.keys()
+    }
+    averaged = {field: total / len(tensor_statistics) for field, total in averaged.items()}
+
+    return TensorStatistics(**averaged)
+
+
+def create_one_hot_observation(vector_length: int, one_hot_index: int | None) -> list[int | float]:
+    """
+    :param vector_length: Length of the one-hot vector.
+    :param one_hot_index: Index of the vector whose element should be set to 1.0, leaving all others as 0.0.
+    :return: Constructed one-hot vector in a list.
+    """
+
+    if one_hot_index is not None and one_hot_index < 0:
+        raise ValueError("One hot indices must be greater than 0.")
+
+    one_hot = np.zeros(vector_length, dtype=np.int8)
+
+    if one_hot_index is not None:
+        one_hot[one_hot_index] = 1
+
+    as_list = one_hot.tolist()
+
+    assert isinstance(as_list, list), "One-hot vector must be a list."
+
+    return as_list
+
+
+def create_one_hot_depth_encoding(agent_controlled_modules: list[str], parameter_group_name: str) -> list[int | float]:
+    """
+    :param agent_controlled_modules: Ordered list of parameter group names in the inner model.
+    :param parameter_group_name: Name of the parameter group to create a depth one-hot vector for.
+    :return: Constructed one-hot depth encoding in a list.
+
+    :note: GANNO encodes depths to one-hot vectors of length 3 regardless of the size of the model.
+    """
+
+    module_index = agent_controlled_modules.index(parameter_group_name)
+    number_of_modules = len(agent_controlled_modules)
+
+    one_hot_index = min(2, (module_index * 3) // number_of_modules)
+
+    return create_one_hot_observation(vector_length=3, one_hot_index=one_hot_index)
+
+
+def tensor_on_local_rank(tensor: torch.Tensor | None) -> bool:
+    """
+    :param tensor: Tensor to check whether it is owned by the local rank, partially or entirely.
+    :return: Whether the tensor is owned by the local rank.
+    """
+
+    return tensor is not None and tensor.grad is not None and tensor.numel() > 0
+
+
+def form_update_tensor(
+    optimizer: optim.Optimizer, parameters: list[torch.Tensor], parameter_group: dict[str, Any]
+) -> None | torch.Tensor:
+    """
+    :param optimizer: Optimizer to form the update tensor from.
+    :param parameters: Parameters to create the update tensor from.
+    :param parameter_group: Parameter group within the optimizer the given parameters came from.
+    :return: None or the formed update tensor.
+    """
+
+    if type(optimizer) in optim_utils.ADAM_OPTIMISERS:
+        updates_list = [optimizer.state[p][EXP_AVERAGE].view(-1) for p in parameters if tensor_on_local_rank(p)]
+        return torch.cat(updates_list) if updates_list else None
+
+    elif type(optimizer) in optim_utils.SGD_OPTIMISERS:
+        return optim_utils.compute_sgd_optimizer_update_stats(
+            optimizer=optimizer, parameter_group=parameter_group, parameters=parameters
+        )
+
+    else:
+        raise NotImplementedError(f"Optimizer {type(optimizer).__name__} is not supported!")
+
+
+def null_standardizer(value_to_standardize: float, **kwargs) -> float:
+    """
+    :param value_to_standardize: Value to mock the standardization of.
+    :return: Given value to standardize.
+    """
+
+    return value_to_standardize
+
+
+def create_sinusoidal_depth_encoding(
+    agent_controlled_modules: list[str], parameter_group_name: str, dimensionality: int
+) -> list[int | float]:
+    """
+    :param agent_controlled_modules: Ordered list of parameter group names in the inner model.
+    :param parameter_group_name: Name of the parameter group to create a depth encoding for.
+    :param dimensionality: Length of the depth vector.
+    :return: Sinusoidal depth encoding.
+    """
+
+    assert dimensionality % 2 == 0, "Dimensionality of a sinusoidal depth encoding must be even."
+
+    depth = agent_controlled_modules.index(parameter_group_name)
+
+    positions = np.arange(dimensionality // 2)
+    frequencies = 1 / (10000 ** (2 * positions / dimensionality))
+
+    encoding = np.zeros(dimensionality)
+    encoding[0::2] = np.sin(depth * frequencies)
+    encoding[1::2] = np.cos(depth * frequencies)
+
+    vector = encoding.tolist()
+
+    return vector
+
+
+def concatenate_lists(lists: list[list[Any]]) -> list[Any]:
+    """
+    :param lists: Lists to concatenate.
+    :return: Concatenated lists.
+    """
+
+    return list(chain(*lists))

libinephany/observations/observer_pipeline.py

@@ -0,0 +1,307 @@
+# ======================================================================================================================
+#
+# IMPORTS
+#
+# ======================================================================================================================
+
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+
+from libinephany.observations.observers.observer_containers import GlobalObserverContainer, LocalObserverContainer
+from libinephany.observations.post_processors import postprocessors
+from libinephany.observations.post_processors.postprocessors import ObservationPostProcessor
+from libinephany.pydantic_models.configs.observer_config import AgentObserverConfig, ObserverConfig
+from libinephany.pydantic_models.schemas.observation_models import ObservationInputs
+from libinephany.pydantic_models.schemas.tensor_statistics import TensorStatistics
+from libinephany.pydantic_models.states.hyperparameter_states import HyperparameterStates
+from libinephany.utils.standardizers import Standardizer
+from libinephany.utils.typing import ObservationInformation
+
+# ======================================================================================================================
+#
+# CLASSES
+#
+# ======================================================================================================================
+
+
+class ObserverPipeline:
+
+    def __init__(
+        self,
+        observer_config: ObserverConfig,
+        agent_config: AgentObserverConfig,
+        standardizer: Standardizer | None,
+        agent_id_to_modules: dict[str, str | None],
+    ):
+        """
+        :param observer_config: ObserverConfig that contains various parameters applicable to all agent observers as
+        well as the observer configs for specific agents.
+        :param agent_config: AgentObserverConfig storing configuration of the agent's actions and observation space for
+        this type of agent.
+        :param standardizer: None or the standardizer to apply to the returned observations.
+        :param agent_id_to_modules: Dictionary mapping agent IDs to None or the name of the parameter group that agent
+        is modulating.
+        """
+
+        self.observer_config = observer_config
+        self.agent_config = agent_config
+        self.standardizer = standardizer
+        self.agent_id_to_modules = agent_id_to_modules
+
+        self.global_observers = GlobalObserverContainer(
+            global_config=observer_config,
+            agent_config=agent_config,
+            standardizer=standardizer,
+        )
+
+        self.local_observers: dict[str, LocalObserverContainer] = self._build_local_observers(
+            agent_id_to_modules=agent_id_to_modules,
+        )
+
+        self.post_processors: list[ObservationPostProcessor] = self._build_post_processors()
+
+        self.clipping_threshold = self.observer_config.observation_clipping_threshold
+        self.invalid_threshold = self.observer_config.invalid_observation_threshold
+
+    @property
+    def observation_size(self) -> int:
+        """
+        :return: Total size of the observation vectors produced by this pipeline.
+        """
+
+        globals_size = self.global_observers.total_observer_size
+        locals_size = list(self.local_observers.values())[0].total_observer_size
+
+        total_size = globals_size + locals_size
+
+        if self.agent_config.prepend_invalid_indicator:
+            total_size += 1
+
+        return total_size
+
+    def _build_local_observers(
+        self,
+        agent_id_to_modules: dict[str, str | None],
+    ) -> dict[str, LocalObserverContainer]:
+        """
+        :param agent_id_to_modules: Dictionary mapping agent IDs to None or the name of the parameter group that agent
+        is modulating.
+        :return: Dictionary mapping agent IDs to LocalObserverContainers for each agent.
+        """
+
+        observers = {}
+
+        for agent_id, parameter_group_name in agent_id_to_modules.items():
+            observers[agent_id] = LocalObserverContainer(
+                global_config=self.observer_config,
+                agent_config=self.agent_config,
+                standardizer=self.standardizer,
+                agent_id=agent_id,
+                parameter_group_name=parameter_group_name,
+            )
+
+        return observers
+
+    def _build_post_processors(self) -> list[ObservationPostProcessor]:
+        """
+        :return: List of post-processors to apply to the global and local observations.
+        """
+
+        if self.agent_config.postprocessors is None:
+            return []
+
+        post_processors = []
+
+        for observer_name, post_processor_kwargs in self.agent_config.postprocessors.items():
+            try:
+                post_processor_type: type[ObservationPostProcessor] = getattr(postprocessors, observer_name)
+
+            except AttributeError as e:
+                raise AttributeError(f"The class {observer_name} does not exist within {postprocessors}!") from e
+
+            if post_processor_kwargs is None:
+                post_processor_kwargs = {}
+
+            post_processor = post_processor_type(
+                observer_config=self.observer_config,
+                **post_processor_kwargs,
+            )
+
+            post_processors.append(post_processor)
+
+        return post_processors
+
+    @staticmethod
+    def merge_globals_to_locals(
+        global_obs: list[float | int],
+        local_obs: dict[str, list[float | int]],
+    ) -> dict[str, list[float | int]]:
+        """
+        :param global_obs: Global observations to post-process.
+        :param local_obs: Dictionary mapping agent IDs to local observations of that agent.
+        :return: Tuple of clipped global and local observations.
+        :return: Dictionary mapping agent ID to that agent's completed observation vector.
+        """
+
+        return {agent_id: global_obs + agent_obs for agent_id, agent_obs in local_obs.items()}
+
+    def get_required_trackers(self) -> list[dict[str, dict[str, Any] | None]]:
+        """
+        :return: List of trackers required by each of the stored observers.
+        """
+
+        required_trackers = self.global_observers.get_required_trackers()
+
+        for agent_container in self.local_observers.values():
+            required_trackers += agent_container.get_required_trackers()
+
+        return required_trackers
+
+    def clip_observation_vector(self, observation_vector: list[float | int]) -> tuple[bool, list[float | int]]:
+        """
+        :param observation_vector: Observations to clip.
+        :return: Tuple indicating whether an observation was clipped and a list of the post-processed observations.
+        """
+
+        invalid_encountered = False
+        post_processed = []
+
+        for observation in observation_vector:
+            if abs(observation) >= self.clipping_threshold:
+                if abs(observation) >= self.invalid_threshold:
+                    invalid_encountered = True
+
+                observation = -self.clipping_threshold if observation < 0 else self.clipping_threshold
+
+            elif np.isnan(observation):
+                observation = self.observer_config.invalid_observation_replacement_value
+                invalid_encountered = True
+
+            post_processed.append(observation)
+
+        return invalid_encountered, post_processed
+
+    def clip_observations(
+        self,
+        globals_to_clip: list[float | int],
+        locals_to_clip: dict[str, list[float | int]],
+    ) -> tuple[list[float | int], dict[str, list[float | int]], bool]:
+        """
+        :param globals_to_clip: Global observations to post-process.
+        :param locals_to_clip: Dictionary mapping agent IDs to local observations of that agent.
+        :return: Tuple of clipped global and local observations and whether any observation clipping occurred.
+        """
+
+        invalid_encountered, globals_to_clip = self.clip_observation_vector(observation_vector=globals_to_clip)
+        post_processed_locals = {}
+
+        for agent_id, agent_observations in locals_to_clip.items():
+            agent_invalid_encountered, post_processed = self.clip_observation_vector(
+                observation_vector=agent_observations
+            )
+
+            post_processed_locals[agent_id] = post_processed
+            invalid_encountered = invalid_encountered if invalid_encountered else agent_invalid_encountered
+
+        if self.agent_config.prepend_invalid_indicator:
+            globals_to_clip = [int(invalid_encountered)] + globals_to_clip
+
+        return globals_to_clip, post_processed_locals, invalid_encountered
+
+    def observe(
+        self,
+        observation_inputs: ObservationInputs,
+        hyperparameter_states: HyperparameterStates,
+        tracked_statistics: dict[str, dict[str, float | TensorStatistics]],
+        actions_taken: dict[str, float | int | None],
+    ) -> tuple[dict[str, list[float | int]], bool]:
+        """
+        :param observation_inputs: Observation input metrics not calculated with statistic trackers.
+        :param hyperparameter_states: HyperparameterStates that manages the hyperparameters.
+        :param tracked_statistics: Dictionary mapping statistic tracker class names to dictionaries mapping module
+        names to floats or TensorStatistic models.
+        :param actions_taken: Dictionary mapping agent IDs to actions taken by that agent.
+        :return: Tuple of a dictionary mapping agent ID to that agent's completed observation vector and a boolean
+        indicating whether an observation clip occurred.
+        """
+
+        global_obs = self.global_observers.observe(
+            observation_inputs=observation_inputs,
+            hyperparameter_states=hyperparameter_states,
+            tracked_statistics=tracked_statistics,
+            action_taken=None,
+        )
+
+        local_obs = {
+            agent_id: agent_observers.observe(
+                observation_inputs=observation_inputs,
+                hyperparameter_states=hyperparameter_states,
+                tracked_statistics=tracked_statistics,
+                action_taken=actions_taken[agent_id] if agent_id in actions_taken else None,
+            )
+            for agent_id, agent_observers in self.local_observers.items()
+        }
+
+        for post_processor in self.post_processors:
+            global_obs, local_obs = post_processor.postprocess(
+                global_observations=global_obs, local_observations=local_obs
+            )
+
+        global_obs, local_obs, obs_clipped = self.clip_observations(
+            globals_to_clip=global_obs, locals_to_clip=local_obs
+        )
+
+        return self.merge_globals_to_locals(global_obs=global_obs, local_obs=local_obs), obs_clipped
+
+    def inform(self) -> tuple[ObservationInformation, dict[str, ObservationInformation]]:
+        """
+        :return: Dictionary of observation info to add to the agent info.
+        """
+
+        global_information = self.global_observers.inform()
+
+        agent_information = {agent_id: observer.inform() for agent_id, observer in self.local_observers.items()}
+
+        return global_information, agent_information
+
+    def get_observation_spaces(self) -> dict[str, gym.spaces.Box]:
+        """
+        :return: Dictionary mapping agent IDs to their observation spaces.
+        """
+
+        return {
+            agent_id: gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.observation_size,), dtype=np.float32)
+            for agent_id in self.local_observers
+        }
+
+    def reset(self) -> None:
+        """
+        Resets all global and local observers.
+        """
+
+        self.global_observers.reset()
+
+        for local_observers in self.local_observers.values():
+            local_observers.reset()
+
+    def train(self) -> None:
+        """
+        Sets all observer containers into training mode.
+        """
+
+        self.global_observers.train()
+
+        for local_observers in self.local_observers.values():
+            local_observers.train()
+
+    def infer(self) -> None:
+        """
+        Sets all observer containers into inference mode.
+        """
+
+        self.global_observers.infer()
+
+        for local_observers in self.local_observers.values():
+            local_observers.infer()