sac.py

import copy
import itertools
from functools import cached_property

import numpy as np
import torch
import torch.nn as nn
from torch import Tensor, optim

from aitraineree import DEVICE
from aitraineree.agents import AgentBase
from aitraineree.agents.agent_utils import hard_update, soft_update
from aitraineree.buffers import PERBuffer
from aitraineree.buffers.buffer_factory import BufferFactory
from aitraineree.loggers import DataLogger
from aitraineree.networks.bodies import ActorBody, CriticBody
from aitraineree.networks.heads import DoubleCritic
from aitraineree.policies import GaussianPolicy, MultivariateGaussianPolicySimple
from aitraineree.types.dataspace import DataSpace
from aitraineree.types.experience import Experience
from aitraineree.types.state import AgentState, BufferState, NetworkState
from aitraineree.utils import to_numbers_seq, to_tensor


class SACAgent(AgentBase):
    """
    Soft Actor-Critic.

    Uses stochastic policy and dual value network (two critics).

    Based on
    "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor"
    by Haarnoja et al. (2018) (http://arxiv.org/abs/1801.01290).

    """

    model = "SAC"

    def __init__(self, obs_space: DataSpace, action_space: DataSpace, **kwargs):
        """
        Parameters:
            obs_space (DataSpace): Dataspace describing the input.
            action_space (DataSpace): Dataspace describing the output.

        Keyword arguments:
            hidden_layers (tuple of ints): Shape of the hidden layers in fully connected network. Default: (128, 128).
            gamma (float):  Discount value. Default: 0.99.
            tau (float): Soft copy fraction. Default: 0.02.
            batch_size (int): Number of samples in a batch. Default: 64.
            buffer_size (int): Size of the prioritized experience replay buffer. Default: 1e6.
            warm_up: (default: 0) Number of samples that needs to be observed before starting to learn. Default: 0.
            update_freq (int): Number of samples between policy updates. Default: 1.
            number_updates (int):  Number of times of batch sampling/training per `update_freq`. Default: 1.
            alpha (float): Weight of log probs in value function. Default: 0.2.
            alpha_lr (Optional float): If not None, it adds alpha as a training parameters with `alpha_lr` as its
                learning rate. Default: None.
            action_scale (float):  Scale for returned action values. Default: 1.
            max_grad_norm_alpha (float): Gradient clipping for the alpha. Default: 1.
            max_grad_norm_actor (float): Gradient clipping for the actor. Default: 10.
            max_grad_norm_critic (float):  Gradient clipping for the critic. Default: 10.
            device: Defaults to CUDA if available. Default: CUDA if available.

        """
        super().__init__(**kwargs)
        self.device = self._register_param(kwargs, "device", DEVICE, update=True)  # Default device is CUDA if available
        self.obs_space = obs_space
        self.action_space = action_space
        self._config["obs_space"] = self.obs_space
        self._config["action_space"] = self.action_space
        action_size = self.action_space.shape[0]  # Because of 1D

        self.gamma: float = float(self._register_param(kwargs, "gamma", 0.99))
        self.tau: float = float(self._register_param(kwargs, "tau", 0.02))
        self.batch_size: int = int(self._register_param(kwargs, "batch_size", 64))
        self.buffer_size: int = int(self._register_param(kwargs, "buffer_size", int(1e6)))
        self.buffer = PERBuffer(self.batch_size, self.buffer_size)

        self.action_scale = self._register_param(kwargs, "action_scale", 1)

        self.warm_up = int(self._register_param(kwargs, "warm_up", 0))
        self.update_freq = int(self._register_param(kwargs, "update_freq", 1))
        self.number_updates = int(self._register_param(kwargs, "number_updates", 1))
        self.critic_number_updates = int(self._register_param(kwargs, "critic_number_updates", 1))
        self.actor_number_updates = int(self._register_param(kwargs, "actor_number_updates", 1))

        # Reason sequence initiation.
        hidden_layers = to_numbers_seq(self._register_param(kwargs, "hidden_layers", (128, 128)))
        actor_hidden_layers = to_numbers_seq(self._register_param(kwargs, "actor_hidden_layers", hidden_layers))
        critic_hidden_layers = to_numbers_seq(self._register_param(kwargs, "critic_hidden_layers", hidden_layers))

        self.simple_policy = bool(self._register_param(kwargs, "simple_policy", True))
        if self.simple_policy:
            self.policy = MultivariateGaussianPolicySimple(action_size, **kwargs)
            self.actor = ActorBody(
                obs_space.shape,
                (self.policy.param_dim * action_size,),
                hidden_layers=actor_hidden_layers,
                gate=nn.ReLU(),
                gate_out=torch.tanh,
                device=self.device,
            )
        else:
            self.policy = GaussianPolicy(
                (actor_hidden_layers[-1],),
                self.action_space.shape,
                out_scale=self.action_scale,
                gate=nn.ReLU(),
                device=self.device,
            )
            self.actor = ActorBody(
                obs_space.shape,
                (actor_hidden_layers[-1],),
                hidden_layers=actor_hidden_layers[:-1],
                gate=nn.ReLU(),
                gate_out=torch.tanh,
                device=self.device,
            )

        self.double_critic = DoubleCritic(
            obs_space.shape,
            action_size,
            CriticBody,
            hidden_layers=critic_hidden_layers,
            gate=nn.ReLU(),
            device=self.device,
        )
        self.target_double_critic = DoubleCritic(
            obs_space.shape,
            action_size,
            CriticBody,
            hidden_layers=critic_hidden_layers,
            gate=nn.ReLU(),
            device=self.device,
        )

        # Target sequence initiation
        hard_update(self.target_double_critic, self.double_critic)

        # Optimization sequence initiation.
        self.target_entropy = -action_size
        alpha_lr = self._register_param(kwargs, "alpha_lr")
        self.alpha_lr = float(alpha_lr) if alpha_lr else None
        alpha_init = float(self._register_param(kwargs, "alpha", 0.2))
        self.log_alpha = torch.tensor(np.log(alpha_init), device=self.device, requires_grad=True)

        actor_lr = float(self._register_param(kwargs, "actor_lr", 3e-4))
        critic_lr = float(self._register_param(kwargs, "critic_lr", 3e-4))

        self.actor_params = list(self.actor.parameters()) + list(self.policy.parameters())
        self.critic_params = list(self.double_critic.parameters())
        self.actor_optimizer = optim.Adam(self.actor_params, lr=actor_lr)
        self.critic_optimizer = optim.Adam(list(self.critic_params), lr=critic_lr)
        if self.alpha_lr is not None:
            self.alpha_optimizer = optim.Adam([self.log_alpha], lr=self.alpha_lr)
        self.max_grad_norm_alpha = float(self._register_param(kwargs, "max_grad_norm_alpha", 1.0))
        self.max_grad_norm_actor = float(self._register_param(kwargs, "max_grad_norm_actor", 10.0))
        self.max_grad_norm_critic = float(self._register_param(kwargs, "max_grad_norm_critic", 10.0))

        # Breath, my child.
        self.iteration = 0

        self._loss_actor = float("nan")
        self._loss_critic = float("nan")
        self._metrics: dict[str, float | dict[str, float]] = {}

    @property
    def alpha(self):
        return self.log_alpha.exp()

    @property
    def loss(self) -> dict[str, float]:
        return {"actor": self._loss_actor, "critic": self._loss_critic}

    @loss.setter
    def loss(self, value):
        if isinstance(value, dict):
            self._loss_actor = value["actor"]
            self._loss_critic = value["critic"]
        else:
            self._loss_actor = value
            self._loss_critic = value

    def __eq__(self, o: object) -> bool:
        return (
            super().__eq__(o)
            and isinstance(o, type(self))
            and self._config == o._config
            and self.buffer == o.buffer
            and self.get_network_state() == o.get_network_state()  # TODO @dawid: Currently net isn't compared properly
        )

    def reset_agent(self) -> None:
        self.actor.reset_parameters()
        self.policy.reset_parameters()
        self.double_critic.reset_parameters()
        hard_update(self.target_double_critic, self.double_critic)

    @cached_property
    def action_min(self):
        return to_tensor(self.action_space.low)

    @cached_property
    def action_max(self):
        return to_tensor(self.action_space.high)

    def state_dict(self) -> dict[str, dict]:
        """
        Returns network's weights in order:
        Actor, TargetActor, Critic, TargetCritic
        """
        return {
            "actor": self.actor.state_dict(),
            "policy": self.policy.state_dict(),
            "double_critic": self.double_critic.state_dict(),
            "target_double_critic": self.target_double_critic.state_dict(),
        }

    @torch.no_grad()
    def act(self, experience: Experience, epsilon: float = 0.0, deterministic: bool = False) -> Experience:
        """Acting on the observations. Returns action.

        Parameters:
            obs (array_like): current state
            eps (float): epsilon, for epsilon-greedy action selection
            deterministic (optional bool): Whether to use deterministic policy. Only has effect in `train` mode.
                In `test` mode all actions are deterministic.

        Returns:
            action: (list float) Action values.

        """
        if self.train and self.iteration < self.warm_up or self._rng.random() < epsilon:
            rnd = torch.rand(self.action_space.shape)
            rnd_action = (self.action_max + self.action_min) * rnd + self.action_min
            action = rnd_action.cpu().tolist()
            return experience.update(action=action)

        _deterministic = (not self.train) or deterministic
        t_obs = to_tensor(experience.obs).view((1,) + self.obs_space.shape).float().to(self.device)

        proto_action = self.actor(t_obs)
        action = self.policy(proto_action, _deterministic).flatten()
        if self.train and not deterministic:
            last_samples = self.policy._last_samples
            added_noise = action - self.policy(proto_action, deterministic=True).flatten()
            noise_params = [*last_samples.get("mu").flatten().tolist(), *last_samples.get("std").flatten().tolist()]
            experience = experience.update(noise=added_noise, noise_params=noise_params)
        action = torch.clamp(action, self.action_min, self.action_max)
        action = action.tolist()
        return experience.update(action=action)

    def step(self, experience: Experience) -> None:
        if not self.train:
            return

        self.iteration += 1
        self.buffer.add(
            obs=experience.obs,
            action=experience.action,
            reward=experience.reward,
            next_obs=experience.next_obs,
            done=experience.done,
        )

        if self.iteration < self.warm_up:
            return

        if len(self.buffer) > self.batch_size and (self.iteration % self.update_freq) == 0:
            for _ in range(self.number_updates):
                self.learn(self.buffer.sample())

    def compute_value_loss(self, states, actions, rewards, next_states, dones) -> tuple[Tensor, Tensor]:
        Q1_expected, Q2_expected = self.double_critic(states, actions)

        with torch.no_grad():
            # next_actions, log_prob = self.actor(next_states)
            proto_next_action = self.actor(next_states)
            next_actions = self.policy(proto_next_action)
            log_prob = self.policy.log_prob(next_actions).view(-1, 1)
            assert next_actions.shape == (self.batch_size,) + self.action_space.shape
            assert log_prob.shape == (self.batch_size, 1)

            Q1_target_next, Q2_target_next = self.target_double_critic.act(next_states, next_actions)
            assert Q1_target_next.shape == Q2_target_next.shape == (self.batch_size, 1)

            Q_min = torch.min(Q1_target_next, Q2_target_next)
            QH_target = Q_min - self.alpha * log_prob
            assert QH_target.shape == (self.batch_size, 1)

            Q_target = rewards + self.gamma * QH_target * (1 - dones)
            assert Q_target.shape == (self.batch_size, 1)

        Q1_diff = Q1_expected - Q_target
        error_1 = Q1_diff.pow(2)
        mse_loss_1: Tensor = error_1.mean()
        self._metrics["value/critic1"] = {"mean": float(Q1_expected.mean()), "std": float(Q1_expected.std())}
        self._metrics["value/critic1_lse"] = float(mse_loss_1.item())

        Q2_diff = Q2_expected - Q_target
        error_2 = Q2_diff.pow(2)
        mse_loss_2: Tensor = error_2.mean()
        self._metrics["value/critic2"] = {"mean": float(Q2_expected.mean()), "std": float(Q2_expected.std())}
        self._metrics["value/critic2_lse"] = float(mse_loss_2.item())

        with torch.no_grad():
            Q_diff = Q1_expected - Q2_expected
            self._metrics["value/Q_diff"] = {"mean": float(Q_diff.mean()), "std": float(Q_diff.std())}
            error: Tensor = torch.max(error_1, error_2)

        loss = mse_loss_1 + mse_loss_2
        return loss, error

    def compute_policy_loss(self, states):
        self.double_critic.requires_grad_ = False
        proto_actions = self.actor(states)
        pred_actions = self.policy(proto_actions)
        log_prob = self.policy.log_prob(pred_actions)
        assert pred_actions.shape == (self.batch_size,) + self.action_space.shape

        Q_estimate = torch.min(*self.double_critic(states, pred_actions))
        assert Q_estimate.shape == (self.batch_size, 1)

        self._metrics["policy/entropy"] = -float(log_prob.detach().mean())
        loss = (self.alpha * log_prob - Q_estimate).mean()

        # Update alpha
        if self.alpha_lr is not None:
            self.alpha_optimizer.zero_grad()
            # loss_alpha = -(self.alpha * (log_prob + self.target_entropy).detach()).mean()  # CORRECT?
            loss_alpha = -(self.alpha * log_prob.detach()).mean()
            loss_alpha.backward()
            nn.utils.clip_grad_norm_(self.log_alpha, self.max_grad_norm_alpha)
            self.alpha_optimizer.step()

        self.double_critic.requires_grad_ = True
        return loss

    def learn(self, samples):
        """update the critics and actors of all the agents"""
        batch_obs_shape = (self.batch_size,) + self.obs_space.shape
        batch_action_shape = (self.batch_size,) + self.action_space.shape

        rewards = to_tensor(samples["reward"]).float().to(self.device).view(self.batch_size, 1)
        dones = to_tensor(samples["done"]).int().to(self.device).view(self.batch_size, 1)
        obss = to_tensor(samples["obs"]).float().to(self.device).view(batch_obs_shape)
        next_obss = to_tensor(samples["next_obs"]).float().to(self.device).view(batch_obs_shape)
        actions = to_tensor(samples["action"]).to(self.device).view(batch_action_shape)

        self.actor_optimizer.zero_grad()
        policy_loss = self.compute_policy_loss(obss)
        policy_loss.backward()
        nn.utils.clip_grad_norm_(self.actor_params, self.max_grad_norm_actor)
        self.actor_optimizer.step()

        # Critic (value) update
        c_loss = np.zeros(self.critic_number_updates)
        for idx in range(self.critic_number_updates):
            self.critic_optimizer.zero_grad()
            value_loss, error = self.compute_value_loss(obss, actions, rewards, next_obss, dones)
            value_loss.backward()
            # nn.utils.clip_grad_norm_(self.critic_params, self.max_grad_norm_critic)
            self.critic_optimizer.step()
            c_loss[idx] = float(value_loss.item())
        self._loss_critic = c_loss.mean()

        # Actor (policy) update
        a_loss = np.zeros(self.actor_number_updates)
        for idx in range(self.actor_number_updates):
            self.actor_optimizer.zero_grad()
            policy_loss = self.compute_policy_loss(obss)
            policy_loss.backward()
            # nn.utils.clip_grad_norm_(self.actor_params, self.max_grad_norm_actor)
            self.actor_optimizer.step()
            a_loss[idx] = float(policy_loss.item())

        self._loss_actor = a_loss.mean()

        if hasattr(self.buffer, "priority_update"):
            assert any(~torch.isnan(error))
            self.buffer.priority_update(samples["index"], error.abs())

        soft_update(self.target_double_critic, self.double_critic, self.tau)

    def log_metrics(self, data_logger: DataLogger, step: int, full_log: bool = False):
        data_logger.log_value("loss/actor", self._loss_actor, step)
        data_logger.log_value("loss/critic", self._loss_critic, step)
        data_logger.log_value("loss/alpha", self.alpha, step)

        if self.simple_policy:
            policy_params = {str(i): v for i, v in enumerate(itertools.chain.from_iterable(self.policy.parameters()))}
            data_logger.log_values_dict("policy/param", policy_params, step)

        for name, value in self._metrics.items():
            if isinstance(value, dict):
                data_logger.log_values_dict(name, value, step)
            else:
                data_logger.log_value(name, value, step)

        if full_log:
            # TODO: Add Policy layers
            for idx, layer in enumerate(self.actor.layers):
                if hasattr(layer, "weight"):
                    data_logger.create_histogram(f"policy/layer_weights_{idx}", layer.weight, step)
                if hasattr(layer, "bias") and layer.bias is not None:
                    data_logger.create_histogram(f"policy/layer_bias_{idx}", layer.bias, step)

            for idx, layer in enumerate(self.double_critic.critic_1.layers):
                if hasattr(layer, "weight"):
                    data_logger.create_histogram(f"critic_1/layer_{idx}", layer.weight, step)
                if hasattr(layer, "bias") and layer.bias is not None:
                    data_logger.create_histogram(f"critic_1/layer_bias_{idx}", layer.bias, step)

            for idx, layer in enumerate(self.double_critic.critic_2.layers):
                if hasattr(layer, "weight"):
                    data_logger.create_histogram(f"critic_2/layer_{idx}", layer.weight, step)
                if hasattr(layer, "bias") and layer.bias is not None:
                    data_logger.create_histogram(f"critic_2/layer_bias_{idx}", layer.bias, step)

    def get_state(self) -> AgentState:
        return AgentState(
            model=self.model,
            obs_space=self.obs_space,
            action_space=self.action_space,
            buffer=copy.deepcopy(self.buffer.get_state()),
            network=copy.deepcopy(self.get_network_state()),
            config=self._config,
        )

    def get_network_state(self) -> NetworkState:
        return NetworkState(
            net=dict(
                policy=self.policy.state_dict(),
                actor=self.actor.state_dict(),
                double_critic=self.double_critic.state_dict(),
                target_double_critic=self.target_double_critic.state_dict(),
            )
        )

    def set_buffer(self, buffer_state: BufferState) -> None:
        self.buffer = BufferFactory.from_state(buffer_state)

    def set_network(self, network_state: NetworkState) -> None:
        self.policy.load_state_dict(network_state.net["policy"])
        self.actor.load_state_dict(network_state.net["actor"])
        self.double_critic.load_state_dict(network_state.net["double_critic"])
        self.target_double_critic.load_state_dict(network_state.net["target_double_critic"])

    @staticmethod
    def from_state(state: AgentState) -> AgentBase:
        config = copy.copy(state.config)
        config.update({"obs_space": state.obs_space, "action_space": state.action_space})
        agent = SACAgent(**config)
        if state.network is not None:
            agent.set_network(state.network)
        if state.buffer is not None:
            agent.set_buffer(state.buffer)
        return agent

    def save_state(self, path: str):
        agent_state = self.get_state()
        torch.save(agent_state, path)

    def load_state(self, path: str):
        agent_state = torch.load(path)
        self._config = agent_state.get("config", {})
        self.__dict__.update(**self._config)

        self.actor.load_state_dict(agent_state["actor"])
        self.policy.load_state_dict(agent_state["policy"])
        self.double_critic.load_state_dict(agent_state["double_critic"])
        self.target_double_critic.load_state_dict(agent_state["target_double_critic"])