libinephany 0.18.1__py3-none-any.whl → 0.19.0__py3-none-any.whl

libinephany/observations/observers/local_observers.py

@@ -4,11 +4,13 @@
 #
 # ======================================================================================================================
 
+import math
 from typing import Any
 
 from libinephany.observations import observation_utils, statistic_trackers
-from libinephany.observations.observation_utils import StatisticStorageTypes
+from libinephany.observations.observation_utils import StatisticStorageTypes, compute_cdf_feature
 from libinephany.observations.observers.base_observers import LocalObserver
+from libinephany.observations.observers.global_observers.constants import LHOPT_CONSTANTS
 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
@@ -1069,3 +1071,127 @@ class PercentageDepth(LocalObserver):
         """
 
         return {}
+
+
+class LogOfNoiseScaleObserver(LocalObserver):
+
+    def __init__(
+        self,
+        *,
+        decay_factor: float = LHOPT_CONSTANTS["DEFAULT_DECAY_FACTOR"],
+        time_window: int = LHOPT_CONSTANTS["DEFAULT_TIME_WINDOW"],
+        skip_statistics: list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        :param decay_factor: Decay factor for CDF calculation in [1, 2.5, 5, 10, 20]
+        :param time_window: Number of time steps to consider for CDF calculation
+        :param skip_statistics: Whether to skip the statistics
+        or use the squared first order gradients as approximations in the same way Adam does.
+        :param kwargs: Miscellaneous keyword arguments.
+        """
+
+        super().__init__(**kwargs)
+
+        self.skip_statistics = skip_statistics
+        self.decay_factor = max(0.0, decay_factor)
+        self.time_window = max(1, time_window)
+
+        # Store time series data for CDF calculation
+        self._time_series: list[tuple[float, float]] = []  # (time, value) pairs
+        self._current_time: float = 0.0
+
+    def _get_observation_format(self) -> StatisticStorageTypes:
+        """
+        :return: Format the observation returns data in. Must be one of the enum attributes in the StatisticStorageTypes
+        enumeration class.
+        """
+
+        return StatisticStorageTypes.VECTOR
+
+    @property
+    def can_standardize(self) -> bool:
+        """
+        :return: Whether the observation can be standardized.
+        """
+
+        return False
+
+    @property
+    def can_inform(self) -> bool:
+        """
+        :return: Whether observations from the observer can be used in the agent info dictionary.
+        """
+
+        return False
+
+    def _update_time(self) -> None:
+        """Update the current time counter."""
+        self._current_time += 1.0
+
+    def _compute_cdf_feature(self, value: float) -> float:
+        """
+        Compute CDF feature for the given value.
+        training loss will be added to the time series after this call.
+        :param value: The value to compute CDF feature for
+        :return: CDF feature value
+        """
+        return compute_cdf_feature(value, self._time_series, self.decay_factor, self._current_time, self.time_window)
+
+    @property
+    def vector_length(self) -> int:
+        """
+        :return: Length of the vector returned by this observation if it returns a vector.
+        """
+        return 2  # [log_noise_scale, cdf_feature]
+
+    def _observe(
+        self,
+        observation_inputs: ObservationInputs,
+        hyperparameter_states: HyperparameterStates,
+        tracked_statistics: dict[str, dict[str, float | TensorStatistics]],
+        action_taken: float | int | None,
+    ) -> float | int | list[int | float] | TensorStatistics:
+        """
+        :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 action_taken: Action taken by the agent this class instance is assigned to.
+        :return: Single float/int, list of floats/ints or TensorStatistics model to add to the observation vector.
+        """
+
+        statistics = tracked_statistics[statistic_trackers.LogOfNoiseScaleStatistics.__name__]
+
+        raw_value = list(statistics.values())[0]  # type: ignore[list-item]
+        assert isinstance(raw_value, float), f"Expected float, got {type(raw_value)}"  # to avoid type errors with mypy
+        batch_size = hyperparameter_states.global_hparams.batch_size.external_value
+        learning_rate = hyperparameter_states.parameter_group_hparams[
+            self.parameter_group_name
+        ].learning_rate.external_value
+
+        log_b_over_epsilon = math.log(batch_size / learning_rate)
+
+        log_noise_scale = raw_value + log_b_over_epsilon
+
+        cdf_feature = self._compute_cdf_feature(log_noise_scale)  # type: ignore[arg-type]
+        self._update_time()
+
+        return [log_noise_scale, cdf_feature]  # type: ignore[list-item]
+
+    def get_required_trackers(self) -> dict[str, dict[str, Any] | None]:
+        """
+        :return: Dictionary mapping statistic tracker class names to kwargs for the class or None if no kwargs are
+        needed.
+        """
+
+        return {
+            statistic_trackers.LogOfNoiseScaleStatistics.__name__: dict(
+                skip_statistics=self.skip_statistics, sample_frequency=LHOPT_CONSTANTS["DEFAULT_SAMPLE_FREQUENCY"]
+            )
+        }
+
+    def reset(self) -> None:
+        """Reset the observer by clearing the time series."""
+        self._time_series = []
+        self._current_time = 0.0