agents/bc_w.py

# Copyright 2022 Twitter, Inc.
# SPDX-License-Identifier: Apache-2.0

import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from utils.logger import logger

from torch.distributions import Distribution, Normal

EPS = 1e-8

LOG_SIG_MAX = 2
LOG_SIG_MIN = -20


class TanhNormal(Distribution):
    """
    Represent distribution of X where
        X ~ tanh(Z)
        Z ~ N(mean, std)

    Note: this is not very numerically stable.
    """
    def __init__(self, normal_mean, normal_std, device, epsilon=1e-6):
        """
        :param normal_mean: Mean of the normal distribution
        :param normal_std: Std of the normal distribution
        :param epsilon: Numerical stability epsilon when computing log-prob.
        """
        self.normal_mean = normal_mean
        self.normal_std = normal_std
        self.normal = Normal(normal_mean, normal_std)
        self.epsilon = epsilon
        self.device = device

    def sample_n(self, n, return_pre_tanh_value=False):
        z = self.normal.sample_n(n)
        if return_pre_tanh_value:
            return torch.tanh(z), z
        else:
            return torch.tanh(z)

    def log_prob(self, value, pre_tanh_value=None):
        """
        :param value: some value, x
        :param pre_tanh_value: arctanh(x)
        :return:
        """
        if pre_tanh_value is None:
            pre_tanh_value = torch.log(
                (1+value) / (1-value+self.epsilon) + self.epsilon
            ) / 2
        return self.normal.log_prob(pre_tanh_value) - torch.log(
            1 - value * value + self.epsilon
        )

    def sample(self, return_pretanh_value=False):
        """
        Sampling without reparameterization.
        """
        z = self.normal.sample().detach()

        if return_pretanh_value:
            return torch.tanh(z), z
        else:
            return torch.tanh(z)

    def rsample(self, return_pretanh_value=False):
        """
        Sampling in the reparameterization case.
        """
        z = (
            self.normal_mean +
            self.normal_std *
            Normal(
                torch.zeros(self.normal_mean.size(), device=self.device),
                torch.ones(self.normal_std.size(), device=self.device)
            ).sample()
        )
        z.requires_grad_()

        if return_pretanh_value:
            return torch.tanh(z), z
        else:
            return torch.tanh(z)


# Implicit Policy
class Actor(nn.Module):
    """
    Gaussian Policy
    """
    def __init__(self,
                 state_dim,
                 action_dim,
                 max_action,
                 device,
                 hidden_sizes=[256, 256],
                 layer_norm=False):
        super(Actor, self).__init__()

        self.layer_norm = layer_norm
        self.base_fc = []
        last_size = state_dim
        for next_size in hidden_sizes:
            self.base_fc += [
                nn.Linear(last_size, next_size),
                nn.LayerNorm(next_size) if layer_norm else nn.Identity(),
                nn.ReLU(inplace=True),
            ]
            last_size = next_size
        self.base_fc = nn.Sequential(*self.base_fc)

        last_hidden_size = hidden_sizes[-1]
        self.last_fc_mean = nn.Linear(last_hidden_size, action_dim)
        self.last_fc_log_std = nn.Linear(last_hidden_size, action_dim)

        self.max_action = max_action
        self.device = device

    def forward(self, state):

        h = self.base_fc(state)
        mean = self.last_fc_mean(h)
        std = self.last_fc_log_std(h).clamp(LOG_SIG_MIN, LOG_SIG_MAX).exp()

        tanh_normal = TanhNormal(mean, std, self.device)
        action, pre_tanh_value = tanh_normal.rsample(return_pretanh_value=True)
        log_prob = tanh_normal.log_prob(action, pre_tanh_value=pre_tanh_value)
        log_prob = log_prob.sum(dim=1, keepdim=True)

        action = action * self.max_action

        return action, log_prob

    def log_prob(self, state, action):
        h = self.base_fc(state)
        mean = self.last_fc_mean(h)
        std = self.last_fc_log_std(h).clamp(LOG_SIG_MIN, LOG_SIG_MAX).exp()

        tanh_normal = TanhNormal(mean, std, self.device)
        log_prob = tanh_normal.log_prob(action)
        log_prob = log_prob.sum(dim=1, keepdim=True)

        return log_prob

    def sample(self,
               state,
               reparameterize=False,
               deterministic=False):

        h = self.base_fc(state)
        mean = self.last_fc_mean(h)
        std = self.last_fc_log_std(h).clamp(LOG_SIG_MIN, LOG_SIG_MAX).exp()

        if deterministic:
            action = torch.tanh(mean) * self.max_action
        else:
            tanh_normal = TanhNormal(mean, std, self.device)
            if reparameterize:
                action = tanh_normal.rsample()
            else:
                action = tanh_normal.sample()
            action = action * self.max_action

        return action


class Critic(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=300):
        super(Critic, self).__init__()
        self.model = nn.Sequential(nn.Linear(state_dim + action_dim, hidden_dim),
                                   nn.ReLU(),
                                   nn.Linear(hidden_dim, hidden_dim),
                                   nn.ReLU(),
                                   nn.Linear(hidden_dim, 1))

    def forward(self, state, action):
        x = torch.cat([state, action], dim=-1)
        return self.model(x)


class BC_W(object):
    def __init__(self,
                 state_dim,
                 action_dim,
                 max_action,
                 device,
                 discount,
                 tau,
                 lr=3e-4,
                 hidden_sizes=[256,256],
                 w_gamma=5.0,
                 c_iter=3,
                 ):

        self.actor = Actor(state_dim, action_dim, max_action,
                           device=device,
                           hidden_sizes=hidden_sizes,
                           layer_norm=False).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)

        self.critic = Critic(state_dim, action_dim).to(device)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-4)

        self.max_action = max_action
        self.action_dim = action_dim
        self.discount = discount
        self.tau = tau
        self.device = device

        self.w_gamma = w_gamma
        self.c_iter = c_iter

    def sample_action(self, state):
        with torch.no_grad():
            state = torch.FloatTensor(state.reshape(1, -1)).to(self.device)
            action = self.actor.sample(state, deterministic=True)
        return action.cpu().data.numpy().flatten()

    def optimize_c(self, state, action_b):
        action_pi = self.actor.sample(state).detach()

        batch_size = state.shape[0]
        alpha = torch.rand((batch_size, 1)).to(self.device)
        a_intpl = (action_pi + alpha * (action_b - action_pi)).requires_grad_(True)
        grads = torch.autograd.grad(outputs=self.critic(state, a_intpl).mean(), inputs=a_intpl, create_graph=True,
                                    only_inputs=True)[0]
        slope = (grads.square().sum(dim=-1) + EPS).sqrt()
        gradient_penalty = torch.max(slope - 1.0, torch.zeros_like(slope)).square().mean()

        logits_p = self.critic(state, action_pi)
        logits_b = self.critic(state, action_b)
        logits_diff = logits_p - logits_b
        critic_loss = - logits_diff.mean() + gradient_penalty * self.w_gamma

        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        return critic_loss.item()

    def optimize_p(self, state, action_b):
        action_pi = self.actor.sample(state)
        logits_p = self.critic(state, action_pi)
        logits_b = self.critic(state, action_b)
        logits_diff = logits_p - logits_b

        # Actor Training
        actor_loss = logits_diff.mean()
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        return actor_loss.item()


    def train(self, replay_buffer, iterations, batch_size=100):

        for it in range(iterations):
            # Sample replay buffer / batch
            state, action, _, _, _ = replay_buffer.sample(batch_size)

            critic_loss = self.optimize_c(state, action)
            actor_loss = self.optimize_p(state, action)

            for _ in range(self.c_iter - 1):
                state, action, _, _, _ = replay_buffer.sample(batch_size)
                critic_loss = self.optimize_c(state, action)

        logger.record_tabular('Actor Loss', actor_loss)
        logger.record_tabular('Critic Loss', critic_loss)

    def save_model(self, dir):
        torch.save(self.actor.state_dict(), f'{dir}/actor.pth')

    def load_model(self, dir):
        self.actor.load_state_dict(torch.load(f'{dir}/actor.pth'))