a2c.py

"""Synchronous advantage actor-critic learner and training entry point."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch
import torch.nn as nn
from gymnasium.vector import VectorEnv
from torch.distributions import Distribution

from deeprl.algorithms.on_policy import (
    OnPolicyConfig,
    OnPolicyLearner,
    train_on_policy,
)
from deeprl.buffers import RolloutBuffer
from deeprl.logger import Logger, MetricHistory
from deeprl.stats import TrainingStats
from deeprl.utils import explained_variance

if TYPE_CHECKING:
    from torch.utils.tensorboard import SummaryWriter

    from deeprl.evaluate import Evaluator


@dataclass(frozen=True)
class A2CConfig(OnPolicyConfig):
    """Optimization hyperparameters specific to synchronous A2C."""

    # Optimization
    lr_actor            :   float        = 3e-4
    lr_critic           :   float        = 1e-3
    ent_coef            :   float        = 0.0      # Entropy-bonus coefficient.
    max_grad_norm       :   float        = 0.5      # Global L2 gradient clipping threshold.
    normalize_advantage :   bool         = False


class A2CLearner(OnPolicyLearner):
    """A2C optimizers and algorithm-specific full-rollout update."""

    def __init__(
        self,
        policy      :   nn.Module,
        value_fn    :   nn.Module,
        cfg         :   A2CConfig,
    ) -> None:
        super().__init__(policy, value_fn, cfg)
        self.cfg: A2CConfig = cfg
        self.optim_actor = torch.optim.SGD(
            self.policy.parameters(),
            lr=self.cfg.lr_actor,
        )
        self.optim_critic = torch.optim.SGD(
            self.value_fn.parameters(),
            lr=self.cfg.lr_critic,
        )

    def update(
        self,
        buffer          :   RolloutBuffer,
        advantages      :   torch.Tensor,
        value_targets   :   torch.Tensor,
    ) -> TrainingStats:
        """Perform one full-rollout A2C update."""

        # ------------------------------------------------------------------
        # 0. Setup
        # ------------------------------------------------------------------

        # Restore module training mode after eval-mode rollout collection. The
        # forward passes below run outside no_grad() and build new autograd graphs.
        self.policy.train()
        self.value_fn.train()

        # Accumulate metrics over the complete update. A2C currently produces
        # one full-rollout batch, B = T * N.
        metric_sums: dict[str, float] = {
            "losses/policy_loss": 0.0,
            "losses/value_loss": 0.0,
            "diagnostics/entropy": 0.0,
            "grads/policy_norm": 0.0,
            "grads/value_norm": 0.0,
        }
        num_samples = 0
        gradient_steps = 0

        # Normalize once over the complete rollout so the update uses one fixed
        # set of advantage weights, matching the other on-policy learners.
        if self.cfg.normalize_advantage and advantages.numel() > 1:
            advantages = (
                advantages - advantages.mean()
            ) / (advantages.std(unbiased=False) + 1e-8)

        for batch in buffer.get(
            advantages=advantages,
            value_targets=value_targets,
            batch_size=None,
        ):

            # --------------------------------------------------------------
            # 1. Rollout re-evaluation
            # --------------------------------------------------------------

            # batch_size=None yields the complete flattened rollout, B = T * N.
            # Observations: (B, *obs_shape); actions: (B, *action_shape).
            action_distributions: Distribution = self.policy(batch.observations)

            # Recompute log pi(A|S) under the current policy. Shape: (B,).
            log_probs = action_distributions.log_prob(batch.actions)

            # Recompute V(S) with gradients enabled. Shape: (B,).
            values: torch.Tensor = self.value_fn(batch.observations)

            # --------------------------------------------------------------
            # 2. Loss computation
            # --------------------------------------------------------------

            # value_loss = 1/2 * mean[(value_target - V(S))^2]
            #
            # value_targets were computed without gradients, so gradients flow
            # only through V(S), giving the critic its semi-gradient update.
            value_loss = 0.5 * (batch.value_targets - values).square().mean()

            # Mean policy entropy across the rollout. This remains differentiable
            # because it contributes to the policy loss when ent_coef is nonzero.
            entropy = action_distributions.entropy().mean()

            # Advantages were computed without gradients and therefore remain
            # fixed weights in the policy-gradient objective.
            batch_advantages = batch.advantages

            # entropy = mean[H(pi(.|S))]
            # policy_loss = -mean[advantage * log pi(A|S)] - ent_coef * entropy
            # The leading minus converts gradient ascent into optimizer descent.
            policy_loss = -(
                (batch_advantages * log_probs).mean()
                + self.cfg.ent_coef * entropy
            )

            # --------------------------------------------------------------
            # 3. Parameter updates
            # --------------------------------------------------------------

            # Critic semi-gradient update with global L2 gradient clipping.
            self.optim_critic.zero_grad(set_to_none=True)
            value_loss.backward()
            value_grad_norm = torch.nn.utils.clip_grad_norm_(
                self.value_fn.parameters(),
                max_norm=self.cfg.max_grad_norm,
            )
            self.optim_critic.step()

            # Policy-gradient update with global L2 gradient clipping.
            self.optim_actor.zero_grad(set_to_none=True)
            policy_loss.backward()
            policy_grad_norm = torch.nn.utils.clip_grad_norm_(
                self.policy.parameters(),
                max_norm=self.cfg.max_grad_norm,
            )
            self.optim_actor.step()

            # --------------------------------------------------------------
            # 4. Batch metric accumulation
            # --------------------------------------------------------------

            # Weight batch metrics by their number of rollout samples.
            batch_samples = batch.observations.shape[0]
            num_samples += batch_samples
            gradient_steps += 1
            for name, value in (
                ("losses/policy_loss", policy_loss),
                ("losses/value_loss", value_loss),
                ("diagnostics/entropy", entropy),
                ("grads/policy_norm", policy_grad_norm),
                ("grads/value_norm", value_grad_norm),
            ):
                metric_sums[name] += value.detach().item() * batch_samples

        # ------------------------------------------------------------------
        # 5. Update metric reduction
        # ------------------------------------------------------------------

        # Produce one mean for the complete update, not one value per batch.
        metrics = {
            name: total / num_samples
            for name, total in metric_sums.items()
        }

        # Explained variance compares collection-time V(S) against the fixed
        # TD(lambda) value targets over the complete rollout.
        metrics.update(
            {
                "diagnostics/explained_variance": explained_variance(
                    buffer.values.flatten(), value_targets.flatten()
                ),
                "charts/policy_learning_rate": self.optim_actor.param_groups[0]["lr"],
                "charts/value_learning_rate": self.optim_critic.param_groups[0]["lr"],
            }
        )

        return TrainingStats(metrics=metrics, gradient_steps=gradient_steps)


def a2c(
    envs            :   VectorEnv,
    policy          :   nn.Module,
    value_fn        :   nn.Module,
    cfg             :   A2CConfig,
    writer          :   SummaryWriter | None = None,
    logger          :   Logger | None = None,
    evaluator       :   Evaluator | None = None,
) -> MetricHistory:
    """Train policy and value networks using synchronous A2C.

    Collects T steps from N environments and performs one full-batch A2C
    update over each rollout through the shared on-policy training engine.
    """

    learner = A2CLearner(policy, value_fn, cfg)

    return train_on_policy(
        envs=envs,
        learner=learner,
        writer=writer,
        logger=logger,
        evaluator=evaluator,
    )