stable-baselines3/torchy_baselines/sac/policies.py

import torch as th
import torch.nn as nn

from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution

# CAP the standard deviation of the actor
LOG_STD_MAX = 2
LOG_STD_MIN = -20


class Actor(BaseNetwork):
    """
    Actor network (policy) for SAC.

    :param obs_dim: (int) Dimension of the observation
    :param action_dim: (int) Dimension of the action space
    :param net_arch: ([int]) Network architecture
    :param activation_fn: (nn.Module) Activation function
    :param use_sde: (bool) Whether to use State Dependent Exploration or not
    :param log_std_init: (float) Initial value for the log standard deviation
    :param full_std: (bool) Whether to use (n_features x n_actions) parameters
        for the std instead of only (n_features,) when using SDE.
    :param sde_net_arch: ([int]) Network architecture for extracting features
        when using SDE. If None, the latent features from the policy will be used.
        Pass an empty list to use the states as features.
    """
    def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
                 use_sde=False, log_std_init=-3, full_std=True, sde_net_arch=None):
        super(Actor, self).__init__()

        latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
        self.latent_pi = nn.Sequential(*latent_pi_net)
        self.use_sde = use_sde
        self.sde_feature_extractor = None

        if self.use_sde:
            latent_sde_dim = net_arch[-1]
            # Separate feature extractor for SDE
            if sde_net_arch is not None:
                self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)

            # TODO: check for the learn_features
            self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
                                                               learn_features=True, squash_output=True)
            self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
                                                                            latent_sde_dim=latent_sde_dim,
                                                                            log_std_init=log_std_init)
        else:
            self.action_dist = SquashedDiagGaussianDistribution(action_dim)
            self.mu = nn.Linear(net_arch[-1], action_dim)
            self.log_std = nn.Linear(net_arch[-1], action_dim)

    def get_std(self):
        """
        Retrieve the standard deviation of the action distribution.
        Only useful when using SDE.
        It corresponds to `th.exp(log_std)` in the normal case,
        but is slightly different when using `expln` function
        (cf StateDependentNoiseDistribution doc).

        :return: (th.Tensor)
        """
        return self.action_dist.get_std(self.log_std)

    def reset_noise(self, batch_size=1):
        """
        Sample new weights for the exploration matrix, when using SDE.

        :param batch_size: (int)
        """
        self.action_dist.sample_weights(self.log_std, batch_size=batch_size)

    def _get_latent(self, obs):
        latent_pi = self.latent_pi(obs)

        if self.sde_feature_extractor is not None:
            latent_sde = self.sde_feature_extractor(obs)
        else:
            latent_sde = latent_pi
        return latent_pi, latent_sde

    def get_action_dist_params(self, obs):
        latent_pi, latent_sde = self._get_latent(obs)

        if self.use_sde:
            mean_actions, log_std = self.mu(latent_pi), self.log_std
        else:
            mean_actions, log_std = self.mu(latent_pi), self.log_std(latent_pi)
            # Original Implementation to cap the standard deviation
            log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
        return mean_actions, log_std, latent_sde

    def forward(self, obs, deterministic=False):
        mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
        if self.use_sde:
            # Note the action is squashed
            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde, deterministic=deterministic)
        else:
            # Note the action is squashed
            action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
        return action

    def action_log_prob(self, obs):
        mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)

        if self.use_sde:
            action, log_prob = self.action_dist.log_prob_from_params(mean_actions, self.log_std, latent_sde)
        else:
            action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
        return action, log_prob


class Critic(BaseNetwork):
    """
    Critic network (q-value function) for SAC.

    :param obs_dim: (int) Dimension of the observation
    :param action_dim: (int) Dimension of the action space
    :param net_arch: ([int]) Network architecture
    :param activation_fn: (nn.Module) Activation function
    """
    def __init__(self, obs_dim, action_dim,
                 net_arch, activation_fn=nn.ReLU):
        super(Critic, self).__init__()

        q1_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
        self.q1_net = nn.Sequential(*q1_net)

        q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
        self.q2_net = nn.Sequential(*q2_net)

        self.q_networks = [self.q1_net, self.q2_net]

    def forward(self, obs, action):
        qvalue_input = th.cat([obs, action], dim=1)
        return [q_net(qvalue_input) for q_net in self.q_networks]

    def q1_forward(self, obs, action):
        return self.q_networks[0](th.cat([obs, action], dim=1))


class SACPolicy(BasePolicy):
    """
    Policy class (with both actor and critic) for SAC.

    :param observation_space: (gym.spaces.Space) Observation space
    :param action_dim: (gym.spaces.Space) Action space
    :param learning_rate: (callable) Learning rate schedule (could be constant)
    :param net_arch: ([int or dict]) The specification of the policy and value networks.
    :param device: (str or th.device) Device on which the code should run.
    :param activation_fn: (nn.Module) Activation function
    :param use_sde: (bool) Whether to use State Dependent Exploration or not
    :param log_std_init: (float) Initial value for the log standard deviation
    :param sde_net_arch: ([int]) Network architecture for extracting features
        when using SDE. If None, the latent features from the policy will be used.
        Pass an empty list to use the states as features.
    """
    def __init__(self, observation_space, action_space,
                 learning_rate, net_arch=None, device='cpu',
                 activation_fn=nn.ReLU, use_sde=False,
                 log_std_init=-3, sde_net_arch=None):
        super(SACPolicy, self).__init__(observation_space, action_space, device)

        if net_arch is None:
            net_arch = [256, 256]

        self.obs_dim = self.observation_space.shape[0]
        self.action_dim = self.action_space.shape[0]
        self.net_arch = net_arch
        self.activation_fn = activation_fn
        self.net_args = {
            'obs_dim': self.obs_dim,
            'action_dim': self.action_dim,
            'net_arch': self.net_arch,
            'activation_fn': self.activation_fn
        }
        self.actor_kwargs = self.net_args.copy()
        sde_kwargs = {
            'use_sde': use_sde,
            'log_std_init': log_std_init,
            'sde_net_arch': sde_net_arch
        }
        self.actor_kwargs.update(sde_kwargs)
        self.actor, self.actor_target = None, None
        self.critic, self.critic_target = None, None

        self._build(learning_rate)

    def _build(self, learning_rate):
        self.actor = self.make_actor()
        self.actor.optimizer = th.optim.Adam(self.actor.parameters(), lr=learning_rate(1))

        self.critic = self.make_critic()
        self.critic_target = self.make_critic()
        self.critic_target.load_state_dict(self.critic.state_dict())
        self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))

    def make_actor(self):
        return Actor(**self.actor_kwargs).to(self.device)

    def make_critic(self):
        return Critic(**self.net_args).to(self.device)


MlpPolicy = SACPolicy

register_policy("MlpPolicy", MlpPolicy)