Source code for rl_mind.nn

"""Neural network helpers.

- `build_mlp` builds a multi-layer perceptron from a list of layer sizes.
- `soft_update` performs the Polyak averaging update of a target network,
  :math:`\\theta' \\leftarrow \\tau \\theta + (1 - \\tau) \\theta'`.
"""

import torch
import torch.nn as nn

__all__ = ["build_mlp", "soft_update"]




def build_mlp(
    sizes: list[int],
    activation: nn.Module = nn.ReLU(),
    output_activation: nn.Module | None = None,
) -> nn.Sequential:
    """Build a multi-layer perceptron

    :param sizes: The sizes of the layers, e.g. `[obs_dim, 64, 64, action_dim]`
    :param activation: The activation module between hidden layers
    :param output_activation: An optional activation after the last layer
        (e.g. `nn.Tanh()` for actions in $[-1, 1]$)
    :returns: A `torch.nn.Sequential` of `len(sizes) - 1` linear layers,
        separated by `activation`
    """
    layers: list[nn.Module] = []
    for i, (n_in, n_out) in enumerate(zip(sizes[:-1], sizes[1:])):
        layers.append(nn.Linear(n_in, n_out))
        if i < len(sizes) - 2:
            layers.append(activation)
        elif output_activation is not None:
            layers.append(output_activation)
    return nn.Sequential(*layers)





def soft_update(source: nn.Module, target: nn.Module, tau: float) -> None:
    """Update the target network parameters towards the source network

    Performs $\\theta' \\leftarrow \\tau \\theta + (1-\\tau) \\theta'$ for
    each parameter $\\theta'$ of `target`, in place and without tracking
    gradients. A small `tau` makes the target network lag behind, which is
    what stabilizes the bootstrapped targets of DDPG/TD3/SAC.

    :param source: The network being trained
    :param target: The target network, updated in place
    :param tau: The interpolation factor, in $[0, 1]$ ($\\tau = 1$ copies the
        source network)
    """
    with torch.no_grad():
        for param, target_param in zip(source.parameters(), target.parameters()):
            target_param.mul_(1.0 - tau).add_(param, alpha=tau)