utils.py

import torch
from torch import nn
from torch import Tensor
import numpy as np


def normalize_advantages(adv: Tensor) -> Tensor:
    """
    Normalizes the given advantages tensor by subtracting its mean and dividing by its
    standard deviation, ensuring numerical stability by adding a small constant to the
    standard deviation.
    """
    return (adv - adv.mean()) / (adv.std() + 1e-8)


def ppo_clipped_surrogate(
        adv: Tensor,
        ratio: Tensor,
        clip: float = 0.2
) -> Tensor:
    """
    Computes the Proximal Policy Optimization (PPO) clipped surrogate objective function.

    The function calculates the PPO loss using a clipped ratio to ensure a balance between
    policy updates and avoiding over-optimization, which can destabilize training.
    The clipping mechanism limits the change in policy entropy within a fixed range.
    """
    obj_unclipped = -adv * ratio
    obj_clipped = -adv * torch.clamp(ratio, min=1-clip, max=1+clip)
    policy_loss = torch.max(obj_unclipped, obj_clipped).mean()

    return policy_loss


def kl_divergence(
        ratio: Tensor,
        log_ratio
) -> Tensor:
    """
    Computes an approximation of the Kullback-Leibler (KL) divergence
    between two probability distributions. The method is based on a mean-field
    approximation and assumes input tensors representing probabilities or distributions.
    """
    with torch.no_grad():
        approx_kl = ((ratio - 1) - log_ratio).mean()
    return approx_kl


def layer_init(
        layer,
        std=np.sqrt(2),
        bias=0.0
):
    """
    Initializes the parameters of a given neural network layer with specific values based on
    given standard deviation and bias.

    This function adjusts the initial values of the weights and bias of the provided layer to
    improve the stability of training and convergence. It uses the provided standard deviation
    and bias value to modify these parameters.

    :param layer: The neural network layer whose parameters need to be initialized.
       Must have weight and bias attributes that can be adjusted.
       The 'weight' is a multi-dimensional tensor representing the layer's weights.
    :param std: A floating-point value specifying the standard deviation to scale the
       layer's weights. It defaults to the square root of 2.
    :param bias: A floating-point value specifying the value to initialize the layer's
       bias term. Defaults to 0.0.
    :return: The layer with its weights and bias initialized based on the given standard
       deviation and bias.
    """
    nn.init.orthogonal_(layer.weight, std)
    nn.init.constant_(layer.bias, bias)

    return layer