libinephany 0.18.1__tar.gz → 0.19.0__tar.gz

libinephany-0.19.0/CODE_VERSION.cfg

@@ -0,0 +1 @@
+0.19.0

{libinephany-0.18.1/libinephany.egg-info → libinephany-0.19.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: libinephany
-Version: 0.18.1
+Version: 0.19.0
 Summary: Inephany library containing code commonly used by multiple subpackages.
 Author-email: Inephany <info@inephany.com>
 License: Apache 2.0

{libinephany-0.18.1 → libinephany-0.19.0}/libinephany/observations/observation_utils.py

@@ -25,6 +25,7 @@ from libinephany.utils import optim_utils
 # ======================================================================================================================
 
 EXP_AVERAGE = "exp_avg"
+MOMENTUM_BUFFER = "momentum_buffer"
 MIN_DECAY_FACTOR = 1e-10
 
 MIN_TOTAL_WEIGHT = 1e-15  # Minimum total weight threshold for numerical stability
@@ -64,10 +65,8 @@ def get_exponential_weighted_average(values: list[int | 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
 
 
@@ -232,6 +231,24 @@ def form_update_tensor(
         raise NotImplementedError(f"Optimizer {type(optimizer).__name__} is not supported!")
 
 
+def form_momentum_tensor(
+    optimizer: optim.Optimizer, parameters: list[torch.Tensor], parameter_group: dict[str, Any]
+) -> None | torch.Tensor:
+    """
+    :param optimizer: Optimizer to form the momentum tensor from.
+    :param parameters: Parameters to create the momentum tensor from.
+    :param parameter_group: Parameter group within the optimizer the given parameters came from.
+    """
+    if type(optimizer) in optim_utils.ADAM_OPTIMISERS:
+        momentum_list = [optimizer.state[p][EXP_AVERAGE].view(-1) for p in parameters if tensor_on_local_rank(p)]
+        return torch.cat(momentum_list) if momentum_list else None
+    elif type(optimizer) in optim_utils.SGD_OPTIMISERS:
+        momentum_list = [optimizer.state[p][MOMENTUM_BUFFER].view(-1) for p in parameters if tensor_on_local_rank(p)]
+        return torch.cat(momentum_list) if momentum_list else None
+    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.

{libinephany-0.18.1 → libinephany-0.19.0}/libinephany/observations/observers/global_observers/__init__.py

@@ -8,7 +8,15 @@
 # ======================================================================================================================
 
 
-from .gradient_observers import GlobalFirstOrderGradients, GlobalSecondOrderGradients, LHOPTGradientVarianceFraction
+from .gradient_observers import (
+    CosineSimilarityObserverOfGradientAndMomentum,
+    CosineSimilarityObserverOfGradientAndUpdate,
+    CosineSimilarityOfGradientAndParameter,
+    GlobalFirstOrderGradients,
+    GlobalSecondOrderGradients,
+    LHOPTGradientVarianceFraction,
+    LHOPTMomentumGradientRatio,
+)
 from .hyperparameter_observers import (
     InitialHyperparameters,
     LHOPTHyperparameterRatio,
@@ -31,8 +39,11 @@ from .model_observers import (
     GlobalLAMBTrustRatio,
     GlobalParameters,
     GlobalParameterUpdates,
+    LHOPTAverageParameterUpdateMagnitudeObserver,
+    LHOPTGlobalLAMBTrustRatio,
     LogRatioOfPreviousAndCurrentParamNormEnvStepObserver,
     LogRatioOfUpdateAndPreviousParamNormEnvStepObserver,
+    LogRatioOfUpdateAndPreviousParamNormInnerStepObserver,
     NumberOfLayers,
     NumberOfParameters,
 )
@@ -51,14 +62,17 @@ __all__ = [
     GlobalFirstOrderGradients.__name__,
     GlobalSecondOrderGradients.__name__,
     LHOPTGradientVarianceFraction.__name__,
+    LHOPTMomentumGradientRatio.__name__,
     GlobalActivations.__name__,
     GlobalParameterUpdates.__name__,
     GlobalParameters.__name__,
     GlobalLAMBTrustRatio.__name__,
     NumberOfParameters.__name__,
     NumberOfLayers.__name__,
+    LHOPTAverageParameterUpdateMagnitudeObserver.__name__,
     LogRatioOfPreviousAndCurrentParamNormEnvStepObserver.__name__,
     LogRatioOfUpdateAndPreviousParamNormEnvStepObserver.__name__,
+    LogRatioOfUpdateAndPreviousParamNormInnerStepObserver.__name__,
     TrainingProgress.__name__,
     EpochsCompleted.__name__,
     ProgressAtEachCheckpoint.__name__,
@@ -66,4 +80,8 @@ __all__ = [
     LHOPTValidationLoss.__name__,
     LHOPTLossRatio.__name__,
     PercentileOfLossAtEachCheckpoint.__name__,
+    LHOPTGlobalLAMBTrustRatio.__name__,
+    CosineSimilarityObserverOfGradientAndMomentum.__name__,
+    CosineSimilarityObserverOfGradientAndUpdate.__name__,
+    CosineSimilarityOfGradientAndParameter.__name__,
 ]

{libinephany-0.18.1 → libinephany-0.19.0}/libinephany/observations/observers/global_observers/constants.py

@@ -20,6 +20,7 @@ class LHOPTConstants(TypedDict):
     ZERO_DIVISION_TOLERANCE: float
     DEFAULT_SAMPLE_FREQUENCY: int
     DEFAULT_VARIANCE_THRESHOLD: float
+    DEFAULT_ENV_STEP_SAMPLE_FREQUENCY: int
 
 
 # Create the constants instance
@@ -36,4 +37,5 @@ LHOPT_CONSTANTS: LHOPTConstants = LHOPTConstants(
     ZERO_DIVISION_TOLERANCE=1e-8,
     DEFAULT_SAMPLE_FREQUENCY=4,
     DEFAULT_VARIANCE_THRESHOLD=1e-6,
+    DEFAULT_ENV_STEP_SAMPLE_FREQUENCY=10,
 )

libinephany-0.19.0/libinephany/observations/observers/global_observers/gradient_observers.py

@@ -0,0 +1,511 @@
+# ======================================================================================================================
+#
+# GRADIENT OBSERVERS
+#
+# ======================================================================================================================
+
+import math
+from typing import Any
+
+from libinephany.observations import observation_utils, statistic_trackers
+from libinephany.observations.observation_utils import StatisticStorageTypes
+from libinephany.observations.observers.base_observers import GlobalObserver
+from libinephany.observations.observers.global_observers.base_classes import LHOPTBaseObserver
+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
+
+
+class GlobalFirstOrderGradients(GlobalObserver):
+
+    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.TENSOR_STATISTICS
+
+    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.FirstOrderGradients.__name__]
+
+        return observation_utils.average_tensor_statistics(tensor_statistics=list(statistics.values()))  # type: ignore
+
+    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.FirstOrderGradients.__name__: dict(skip_statistics=self.skip_statistics)}
+
+
+class GlobalSecondOrderGradients(GlobalObserver):
+
+    def __init__(
+        self,
+        *,
+        compute_hessian_diagonal: bool = False,
+        **kwargs,
+    ) -> None:
+        """
+        :param compute_hessian_diagonal: Whether to compute the Hessian diagonal to determine second order gradients
+        or use the squared first order gradients as approximations in the same way Adam does.
+        :param kwargs: Miscellaneous keyword arguments.
+        """
+
+        super().__init__(**kwargs)
+
+        self.compute_hessian_diagonal = compute_hessian_diagonal
+
+    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.TENSOR_STATISTICS
+
+    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.SecondOrderGradients.__name__]
+
+        return observation_utils.average_tensor_statistics(tensor_statistics=list(statistics.values()))  # type: ignore
+
+    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.SecondOrderGradients.__name__: dict(
+                skip_statistics=self.skip_statistics, compute_hessian_diagonal=self.compute_hessian_diagonal
+            )
+        }
+
+
+class LHOPTGradientVarianceFraction(LHOPTBaseObserver):
+    """
+    This is a global observer from the OpenAI paper "Learning to Optimize with Reinforcement Learning"
+    https://arxiv.org/abs/2305.18291.
+
+    It returns two-dimensional observations: [raw_value, cdf_feature] for gradient variance fraction values.
+    """
+
+    def __init__(
+        self,
+        *,
+        variance_threshold: float = LHOPT_CONSTANTS["DEFAULT_VARIANCE_THRESHOLD"],
+        **kwargs,
+    ) -> None:
+        """
+        :param variance_threshold: Threshold for variance comparison in gradient variance fraction calculation
+        :param kwargs: Other observation keyword arguments.
+        """
+        super().__init__(**kwargs)
+        self.variance_threshold = variance_threshold
+
+    @property
+    def can_standardize(self) -> bool:
+        """
+        This observer has its own CDF calculation, no need to standardize.
+        :return: Whether the observation can be standardized.
+        """
+        return False
+
+    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 vector_length(self) -> int:
+        """
+        :return: Length of the vector returned by this observation if it returns a vector.
+        """
+        return 2  # [raw_value, 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.
+        """
+        if statistic_trackers.GradientVarianceFraction.__name__ not in tracked_statistics:
+            return [0.0, 0.0]
+
+        raw_value = list(tracked_statistics[statistic_trackers.GradientVarianceFraction.__name__].values())[0]  # type: ignore[list-item]
+
+        cdf_feature = self._compute_cdf_feature(raw_value)  # type: ignore[arg-type]
+        self._update_time()
+
+        return [raw_value, 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.GradientVarianceFraction.__name__: dict(
+                variance_threshold=self.variance_threshold, sample_frequency=LHOPT_CONSTANTS["DEFAULT_SAMPLE_FREQUENCY"]
+            ),
+        }
+
+
+class LHOPTMomentumGradientRatio(LHOPTBaseObserver):
+    """
+    This is a global observer from the OpenAI paper "Learning to Optimize with Reinforcement Learning"
+    https://arxiv.org/abs/2305.18291.
+
+    It returns two-dimensional observations: [raw_value, cdf_feature] for momentum gradient ratio values.
+    """
+
+    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 vector_length(self) -> int:
+        """
+        :return: Length of the vector returned by this observation if it returns a vector.
+        """
+        return 2  # [raw_value, 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.MomentumGradientRatioStatistics.__name__]
+
+        raw_value = list(statistics.values())[0]  # type: ignore[list-item]
+
+        cdf_feature = self._compute_cdf_feature(raw_value)  # type: ignore[arg-type]
+        self._update_time()
+
+        return [raw_value, 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.MomentumGradientRatioStatistics.__name__: dict(
+                skip_statistics=self.skip_statistics,
+                sample_frequency=LHOPT_CONSTANTS["DEFAULT_SAMPLE_FREQUENCY"],
+            ),
+        }
+
+
+class CosineSimilarityObserverOfGradientAndMomentum(LHOPTBaseObserver):
+    """
+    This is a global observer from the OpenAI paper "Learning to Optimize with Reinforcement Learning"
+    https://arxiv.org/abs/2305.18291.
+
+    It returns two-dimensional observations: [raw_value, cdf_feature] for cosine similarity of gradient and momentum values.
+    """
+
+    def __init__(
+        self,
+        *,
+        skip_statistics: list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        :param compute_hessian_diagonal: Whether to compute the Hessian diagonal to determine second order gradients
+        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
+
+    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 vector_length(self) -> int:
+        """
+        :return: Length of the vector returned by this observation if it returns a vector.
+        """
+        return 3  # [raw_value, cdf_feature, logit_of_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.CosineSimilarityObserverOfGradientAndMomentumStatistics.__name__
+        ]
+
+        raw_value = list(statistics.values())[0]  # type: ignore[list-item]
+
+        cdf_feature = self._compute_cdf_feature(raw_value)  # type: ignore[arg-type]
+        self._update_time()
+
+        # Handle edge cases for logit calculation
+        if cdf_feature <= 0.0 or cdf_feature >= 1.0:
+            logit_of_cdf_feature = 0.0
+        else:
+            logit_of_cdf_feature = math.log(cdf_feature / (1 - cdf_feature))
+
+        return [raw_value, cdf_feature, logit_of_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.CosineSimilarityObserverOfGradientAndMomentumStatistics.__name__: dict(
+                skip_statistics=self.skip_statistics,
+                sample_frequency=LHOPT_CONSTANTS["DEFAULT_SAMPLE_FREQUENCY"],
+            )
+        }
+
+
+class CosineSimilarityObserverOfGradientAndUpdate(LHOPTBaseObserver):
+    """
+    This is a global observer from the OpenAI paper "Learning to Optimize with Reinforcement Learning"
+    https://arxiv.org/abs/2305.18291.
+
+    It returns two-dimensional observations: [raw_value, cdf_feature] for cosine similarity of gradient and update values.
+    """
+
+    def __init__(
+        self,
+        *,
+        skip_statistics: list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        :param skip_statistics: List of statistics to skip.
+        :param kwargs: Miscellaneous keyword arguments.
+        """
+
+        super().__init__(**kwargs)
+
+        self.skip_statistics = skip_statistics
+
+    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 vector_length(self) -> int:
+        """
+        :return: Length of the vector returned by this observation if it returns a vector.
+        """
+        return 3  # [raw_value, cdf_feature, logit_of_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.CosineSimilarityObserverOfGradientAndUpdateStatistics.__name__
+        ]
+
+        raw_value = list(statistics.values())[0]  # type: ignore[list-item]
+
+        cdf_feature = self._compute_cdf_feature(raw_value)  # type: ignore[arg-type]
+        self._update_time()
+
+        # Handle edge cases for logit calculation
+        if cdf_feature <= 0.0 or cdf_feature >= 1.0:
+            logit_of_cdf_feature = 0.0
+        else:
+            logit_of_cdf_feature = math.log(cdf_feature / (1 - cdf_feature))
+
+        return [raw_value, cdf_feature, logit_of_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.CosineSimilarityObserverOfGradientAndUpdateStatistics.__name__: dict(
+                skip_statistics=self.skip_statistics,
+                sample_frequency=LHOPT_CONSTANTS["DEFAULT_SAMPLE_FREQUENCY"],
+            )
+        }
+
+
+class CosineSimilarityOfGradientAndParameter(LHOPTBaseObserver):
+    """
+    This is a global observer from the OpenAI paper "Learning to Optimize with Reinforcement Learning"
+    https://arxiv.org/abs/2305.18291.
+
+    It returns two-dimensional observations: [raw_value, cdf_feature] for cosine similarity of gradient and parameter values.
+    """
+
+    def __init__(
+        self,
+        *,
+        skip_statistics: list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        :param skip_statistics: List of statistics to skip.
+        :param kwargs: Miscellaneous keyword arguments.
+        """
+        super().__init__(**kwargs)
+
+        self.skip_statistics = skip_statistics
+
+    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 vector_length(self) -> int:
+        """
+        :return: Length of the vector returned by this observation if it returns a vector.
+        """
+        return 3  # [raw_value, cdf_feature, logit_of_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.CosineSimilarityOfGradientAndParameterStatistics.__name__]
+
+        raw_value = list(statistics.values())[0]  # type: ignore[list-item]
+
+        cdf_feature = self._compute_cdf_feature(raw_value)  # type: ignore[arg-type]
+        self._update_time()
+
+        # Handle edge cases for logit calculation
+        if cdf_feature <= 0.0 or cdf_feature >= 1.0:
+            logit_of_cdf_feature = 0.0
+        else:
+            logit_of_cdf_feature = math.log(cdf_feature / (1 - cdf_feature))
+
+        return [raw_value, cdf_feature, logit_of_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.CosineSimilarityOfGradientAndParameterStatistics.__name__: dict(
+                skip_statistics=self.skip_statistics,
+                sample_frequency=LHOPT_CONSTANTS["DEFAULT_SAMPLE_FREQUENCY"],
+            )
+        }