Model.py

"""Recurrent actor-critic network."""
from __future__ import annotations

from typing import Optional, Tuple

import torch
import torch.nn as nn
from torch.distributions import Categorical

Hidden = Tuple[torch.Tensor, torch.Tensor]


def init_weights(m: nn.Module, gain: float = 1, bias: float = 0, method: str = "kaiming") -> None:
    """Initialise a layer in place.

    Linear layers: ``kaiming`` (default), ``ortho`` or ``xavier`` weights, constant
    bias. Recurrent layers: Xavier input weights, orthogonal recurrent weights,
    constant bias.
    """
    if isinstance(m, nn.Linear):
        nn.init.constant_(m.bias, bias)
        if method == "kaiming":
            nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
        elif method == "ortho":
            nn.init.orthogonal_(m.weight, gain)
        elif method == "xavier":
            nn.init.xavier_uniform_(m.weight, gain)
        else:
            raise ValueError(f"unknown init method {method!r}")
    elif isinstance(m, (nn.LSTM, nn.RNN, nn.GRU)):
        for name, param in m.named_parameters():
            if "bias" in name:
                nn.init.constant_(param, bias)
            elif "weight_ih" in name:
                nn.init.xavier_uniform_(param)
            elif "weight_hh" in name:
                nn.init.orthogonal_(param)
            else:
                raise ValueError(f"unexpected parameter {name!r}")


class ActorCritic(nn.Module):
    """Two independent recurrent streams: an actor and a critic.

    Parameters
    ----------
    STATE_DIM, ACTION_DIM
        Input and action sizes.
    RNN_SIZE
        Hidden size of each recurrent stream.
    FC_SIZE
        Kept for configuration compatibility; the heads are single linear layers.
    RNN
        Recurrent layer class (``torch.nn.LSTM`` by default in the config).
    """

    def __init__(self, STATE_DIM: int, ACTION_DIM: int, RNN_SIZE: int, FC_SIZE: int, RNN=nn.LSTM) -> None:
        super().__init__()
        self.STATE_DIM = STATE_DIM
        self.ACTION_DIM = ACTION_DIM

        self.rnn1 = RNN(input_size=STATE_DIM, hidden_size=RNN_SIZE)   # actor
        self.l1 = nn.Linear(RNN_SIZE, self.ACTION_DIM)                # policy logits
        self.rnn2 = RNN(input_size=STATE_DIM, hidden_size=RNN_SIZE)   # critic
        self.l2 = nn.Linear(RNN_SIZE, 1)                              # state value

        init_weights(self.rnn1)
        init_weights(self.rnn2)
        init_weights(self.l1, gain=0.01, method="ortho")
        init_weights(self.l2, gain=1, method="ortho")

    def construct_dist(self, x: torch.Tensor) -> Categorical:
        """Categorical policy from actor features."""
        return Categorical(logits=self.l1(x))

    def act(self, x: torch.Tensor, hidden_in: Hidden, enable_noise: bool = True, return_dist: bool = False):
        """Sample (or take the mode of) an action for one time step.

        ``x`` has shape ``(1, batch, STATE_DIM)``. Returns ``action``, its
        log-probability (both ``(1, batch, 1)``), the new hidden state, and
        optionally the distribution.
        """
        x, hidden_out = self.rnn1(x, hidden_in)
        dist = self.construct_dist(x)
        action = dist.sample() if enable_noise else dist.mode
        action_logprob = dist.log_prob(action)
        if return_dist:
            return action.unsqueeze(-1), action_logprob.unsqueeze(-1), hidden_out, dist
        return action.unsqueeze(-1), action_logprob.unsqueeze(-1), hidden_out

    def evaluate(self, x: torch.Tensor, hidden_in: Optional[Hidden] = None):
        """State values ``(T, batch, 1)`` from the critic stream."""
        x, hidden_out = self.rnn2(x) if hidden_in is None else self.rnn2(x, hidden_in)
        return self.l2(x), hidden_out

    def forward(self, x: torch.Tensor, action: torch.Tensor, hidden_in):
        """Re-evaluate a stored rollout for the PPO update.

        Returns log-probabilities of ``action`` under the current policy, state
        values and policy entropy, all ``(T, batch, 1)``. ``hidden_in`` is a pair
        ``[actor_hidden, critic_hidden]`` of initial states.
        """
        a, _ = self.rnn1(x, hidden_in[0])
        dist = self.construct_dist(a)
        action_logprob = dist.log_prob(action.squeeze(-1))
        dist_entropy = dist.entropy()
        v, _ = self.rnn2(x, hidden_in[1])
        v = self.l2(v)
        return action_logprob.unsqueeze(-1), v, dist_entropy.unsqueeze(-1)