ppo.py

import gurobipy as gp
import torch
import torch.nn as nn
from gurobipy import GRB
from torch.distributions import Categorical
from torch.distributions import MultivariateNormal


class RolloutBuffer:
    """
    A simple storage for the rollout buffer
    """
    def __init__(self):
        self.actions = []
        self.states = []
        self.logprobs = []
        self.rewards = []
        self.is_terminals = []

    def clear(self):
        del self.actions[:]
        del self.states[:]
        del self.logprobs[:]
        del self.rewards[:]
        del self.is_terminals[:]
    
    def remove_last(self):
        self.actions.pop()
        self.states.pop()
        self.logprobs.pop()
        self.rewards.pop()
        self.is_terminals.pop()
    
    def __len__(self):
        return len(self.actions)


class ActorCritic(nn.Module):
    def __init__(
        self,
        state_dim,
        action_dim,
        has_continuous_action_space,
        action_std_init,
        device,
        constraint=None,
        temp_bound=(20, 23.5)
    ):
        """
        Initialize the ActorCritic module

        Parameters
        ----------
        state_dim : int
            Number of state dimensions
        action_dim : int
            Number of action dimensions
        has_continuous_action_space : bool
            Whether the action space is continuous or discrete
        action_std_init : float
            Initial standard deviation of the action distribution
        device : torch.device
            Device on which the tensors are allocated
        constraint : dict
            Dictionary containing the weights and biases of the constraint model
        temp_bound : tuple
            Tuple containing the lower and upper bounds of the temperature
        """
        super(ActorCritic, self).__init__()

        self.has_continuous_action_space = has_continuous_action_space
        self.device = device

        if has_continuous_action_space:
            self.action_dim = action_dim
            self.action_var = torch.full((action_dim,), action_std_init * action_std_init).to(self.device)

        # actor
        if has_continuous_action_space:
            self.actor = nn.Sequential(
                nn.Linear(state_dim, 64),
                nn.Tanh(),
                nn.Linear(64, 64),
                nn.Tanh(),
                nn.Linear(64, action_dim),
            )
        else:
            self.actor = nn.Sequential(
                nn.Linear(state_dim, 64),
                nn.Tanh(),
                nn.Linear(64, 64),
                nn.Tanh(),
                nn.Linear(64, action_dim),
                nn.Softmax(dim=-1)
            )

        # critic
        self.critic = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )

        self.constraint = constraint
        self.temp_bound = temp_bound

    def set_action_std(self, new_action_std):
        """
        Set the standard deviation of the action distribution
        Parameters
        ----------
        new_action_std : float
            New standard deviation of the action distribution
        """
        if self.has_continuous_action_space:
            self.action_var = torch.full((self.action_dim,), new_action_std * new_action_std).to(self.device)
        else:
            print("--------------------------------------------------------------------------------------------")
            print("WARNING : Calling ActorCritic::set_action_std() on discrete action space policy")
            print("--------------------------------------------------------------------------------------------")

    def forward(self):
        raise NotImplementedError

    def actor_inference(self, state, predictor_input=None, env=None):
        """
        Perform inference on the actor network, with the constraint model if available.

        Parameters
        ----------
        state : torch.Tensor
            State input to the actor network
        predictor_input : list
            List of inputs to the constraint model
        env : gurobi.Env
            Environment object

        Returns
        -------
            action_mean : torch.Tensor
        """
        action_mean = self.actor(state)
        if self.constraint is None:
            return action_mean

        # Predefined ICNN model design
        model_design = (12, 100, 100, 1)

        # Gurobi optimization
        m = gp.Model(env=env)
        m.Params.LogToConsole = 0
        final_action = m.addMVar(1, lb=0.1, ub=1, vtype='C')
        given_action = m.addMVar(1, lb=0.1, ub=1, vtype='C')
        given_action_distance = m.addMVar(1, lb=0, ub=0.9, vtype='C')
        z = m.addMVar((len(predictor_input) + 1,), lb=-1e9, ub=1e9, vtype='C', name='z')
        m.addConstr(z[:-1] == predictor_input, name="c0")
        m.addConstr(z[-1] == final_action, name="c1")
        m.addConstr(action_mean.detach().cpu().numpy()[0] == given_action, name="c2")

        for ii in range(1, len(model_design)):  # per layer
            zn_raw = m.addMVar((model_design[ii],), lb=-1e7, ub=1e7, vtype='C', name=f"z_raw_{ii}")
            if ii == 1:
                m.addConstr(zn_raw == z @ self.constraint[f"Wzs.{ii - 1}.weight"].detach().numpy().T +
                            self.constraint[f"Wzs.{ii - 1}.bias"].detach().numpy(), name=f"layer_{ii}")
            else:
                if f"Wxs.{ii - 2}.bias" in self.constraint.keys():
                    m.addConstr(zn_raw == zn @ self.constraint[f"Wzs.{ii - 1}.weight"].detach().numpy().T +
                                z @ self.constraint[f"Wxs.{ii - 2}.weight"].detach().numpy().T +
                                self.constraint[f"Wxs.{ii - 2}.bias"].detach().numpy(),
                                name=f"middle layer_{ii}")
                else:
                    m.addConstr(zn_raw == zn @ self.constraint[f"Wzs.{ii - 1}.weight"].detach().numpy().T +
                                z @ self.constraint[f"Wxs.{ii - 2}.weight"].detach().numpy().T,
                                name=f"final_layer_{ii}")

            if ii != len(model_design) - 1:
                zn = m.addMVar((model_design[ii],), lb=0, ub=1e7, vtype='C', name="zn")
                for ij in range(model_design[ii]):
                    m.addGenConstrMax(zn[ij], [zn_raw[ij], 0], name="maximize")

        # Set temperature constraints, assume box constraints between 18 and 26
        m.addConstr(zn_raw >= self.temp_bound[0])
        m.addConstr(zn_raw <= self.temp_bound[1])

        m.addConstr(given_action_distance == (final_action - given_action) * (final_action - given_action))

        # Use g in the objectives function
        m.setObjective(
            given_action_distance,
            GRB.MINIMIZE
        )

        # Optimize and edge case handling
        # m.Params.FeasibilityTol = 1e-4
        m.optimize()

        try:
            print("Optimal action:", final_action.x)
            action_mean[0] = final_action.x
        except gp.GurobiError:
            pass
        return action_mean

    def get_mean(self, state):
        """
        Get the mean of the action distribution
        Parameters
        ----------
        state : torch.Tensor
            State input to the actor network

        Returns
        -------
            action_mean : torch.Tensor
        """
        if self.has_continuous_action_space:
            action_mean = self.actor(state)
            return action_mean.detach()

    def act(self, state):
        """
        Get an action from the actor network

        Parameters
        ----------
        state : torch.Tensor
            State input to the actor network

        Returns
        -------
            action : torch.Tensor
            action_logprob : torch.Tensor
        """
        if self.has_continuous_action_space:
            action_mean = self.actor(state)
            cov_mat = torch.diag(self.action_var).unsqueeze(dim=0)
            dist = MultivariateNormal(action_mean, cov_mat)
        else:
            action_probs = self.actor(state)
            dist = Categorical(action_probs)

        action = dist.sample()
        action_logprob = dist.log_prob(action)

        return action.detach(), action_logprob.detach()

    def evaluate(self, state, action):
        """
        Evaluate the action given state. Used for updating the policy.

        Parameters
        ----------
        state : torch.Tensor
            State input to the actor network
        action : torch.Tensor
            Action input to the actor network

        Returns
        -------
            action_logprobs : torch.Tensor
            state_values : torch.Tensor
            dist_entropy : torch.Tensor
        """
        if self.has_continuous_action_space:
            action_mean = self.actor(state)
            action_var = self.action_var.expand_as(action_mean)
            cov_mat = torch.diag_embed(action_var).to(self.device)
            dist = MultivariateNormal(action_mean, cov_mat)

            # For Single Action Environments.
            if self.action_dim == 1:
                action = action.reshape(-1, self.action_dim)

        else:
            action_probs = self.actor(state)
            dist = Categorical(action_probs)
        action_logprobs = dist.log_prob(action)
        dist_entropy = dist.entropy()
        state_values = self.critic(state)

        return action_logprobs, state_values, dist_entropy


class PPO:
    def __init__(self, state_dim, action_dim, lr_actor, lr_critic, gamma, k_epochs, eps_clip,
                 has_continuous_action_space, action_std_init=0.6, device=torch.device("cpu"),
                 diverse_policies=list(), diverse_weight=0,
                 diverse_weight_alpha=0.99, diverse_increase=True):

        self.has_continuous_action_space = has_continuous_action_space

        if has_continuous_action_space:
            self.action_std = action_std_init

        self.gamma = gamma
        self.eps_clip = eps_clip
        self.k_epochs = k_epochs

        self.buffer = RolloutBuffer()
        self.device = device

        self.policy = ActorCritic(state_dim, action_dim, has_continuous_action_space, action_std_init, device).to(device)
        self.optimizer = torch.optim.Adam([
            {'params': self.policy.actor.parameters(), 'lr': lr_actor},
            {'params': self.policy.critic.parameters(), 'lr': lr_critic}
        ])

        self.policy_old = ActorCritic(state_dim, action_dim, has_continuous_action_space, action_std_init, device).to(device)
        self.policy_old.load_state_dict(self.policy.state_dict())

        self.MseLoss = nn.MSELoss()
        
        self.other_policies = list()
        self.diverse_weight_limit = diverse_weight
        self.diverse_weight_alpha = diverse_weight_alpha
        self.diverse_weight = 0 if diverse_increase else diverse_weight
        self.diverse_weight_update_function = self.increase_weight_function if diverse_increase else self.decrease_weight_function
        for checkpoint_path in diverse_policies:
            other_policy = ActorCritic(state_dim, action_dim, has_continuous_action_space, action_std_init, device).to(device)
            other_policy.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))
            self.other_policies.append(other_policy)
    
    def increase_weight_function(self):
        self.diverse_weight = self.diverse_weight_alpha * self.diverse_weight + self.diverse_weight_limit * (1 - self.diverse_weight_alpha)

    def decrease_weight_function(self):
        self.diverse_weight *= self.diverse_weight_alpha

    def set_action_std(self, new_action_std):

        if self.has_continuous_action_space:
            self.action_std = new_action_std
            self.policy.set_action_std(new_action_std)
            self.policy_old.set_action_std(new_action_std)

        else:
            print("--------------------------------------------------------------------------------------------")
            print("WARNING : Calling PPO::set_action_std() on discrete action space policy")
            print("--------------------------------------------------------------------------------------------")

    def decay_action_std(self, action_std_decay_rate, min_action_std):
        print("--------------------------------------------------------------------------------------------")

        if self.has_continuous_action_space:
            self.action_std = self.action_std - action_std_decay_rate
            self.action_std = round(self.action_std, 4)
            if self.action_std <= min_action_std:
                self.action_std = min_action_std
                print("setting actor output action_std to min_action_std : ", self.action_std)
            else:
                print("setting actor output action_std to : ", self.action_std)
            self.set_action_std(self.action_std)

        else:
            print("WARNING : Calling PPO::decay_action_std() on discrete action space policy")

        print("--------------------------------------------------------------------------------------------")

    def select_action(self, state):
        with torch.no_grad():
            state = torch.FloatTensor(state).to(self.device)
            action, action_logprob = self.policy_old.act(state)

        self.buffer.states.append(state)
        self.buffer.actions.append(action)
        self.buffer.logprobs.append(action_logprob)

        if self.has_continuous_action_space:
            return action.detach().cpu().numpy().flatten()
        else:
            return action.item()

    def learn_action_offline(self, state, action):
        action = [action] if not isinstance(action, list) else action
        with torch.no_grad():
            state = torch.FloatTensor(state).to(self.device)
            action = torch.FloatTensor(action).to(self.device)
            action_logprob, _, _ = self.policy_old.evaluate(state, action)

        self.buffer.states.append(state)
        self.buffer.actions.append(action.detach())
        self.buffer.logprobs.append(action_logprob.detach())

    def get_mean(self, state):
        with torch.no_grad():
            state = torch.FloatTensor(state).to(self.device)
            action = self.policy_old.get_mean(state)
        return action.detach().cpu().numpy().flatten()

    def update(self):
        if len(self.buffer) == 0:
            return
        # Monte Carlo estimate of returns
        rewards = []
        discounted_reward = 0
        for reward, is_terminal in zip(reversed(self.buffer.rewards), reversed(self.buffer.is_terminals)):
            if is_terminal:
                discounted_reward = 0
            discounted_reward = reward + (self.gamma * discounted_reward)
            rewards.insert(0, discounted_reward)

        # Normalizing the rewards
        rewards = torch.tensor(rewards, dtype=torch.float32).to(self.device)
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-7)

        # convert list to tensor
        old_states = torch.squeeze(torch.stack(self.buffer.states, dim=0)).detach().to(self.device)
        old_actions = torch.squeeze(torch.stack(self.buffer.actions, dim=0)).detach().to(self.device)
        old_logprobs = torch.squeeze(torch.stack(self.buffer.logprobs, dim=0)).detach().to(self.device)

        # Optimize policy for K epochs
        for _ in range(self.k_epochs):
            # Evaluating old actions and values
            logprobs, state_values, dist_entropy = self.policy.evaluate(old_states, old_actions)

            # match state_values tensor dimensions with rewards tensor
            state_values = torch.squeeze(state_values)

            # Finding the ratio (pi_theta / pi_theta__old)
            ratios = torch.exp(logprobs - old_logprobs.detach())

            # Finding Surrogate Loss
            advantages = rewards - state_values.detach()
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
            
            # Adding diversity term compared to other policies
            dipg_loss = torch.zeros_like(surr1)
            for policy in self.other_policies:
                # for each other policy, calculate the policy distance ratio and advantages
                # We want to maximize the distance ratio and make advantage as much as possible
                other_logprobs, other_state_values, _ = policy.evaluate(old_states, old_actions)
                other_state_values = torch.squeeze(other_state_values)
                ratios = torch.exp(logprobs - other_logprobs)
                advantages = rewards - other_state_values.detach()
                
                other_surr1 = ratios * advantages
                other_surr2 = torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
                dipg_loss += torch.min(other_surr1, other_surr2)

            # final loss of clipped objective PPO
            if len(self.other_policies):
                dipg_loss /= len(self.other_policies)
            loss = -torch.min(surr1, surr2) + 0.5 * self.MseLoss(state_values, rewards) - 0.01 * dist_entropy + self.diverse_weight * dipg_loss

            # take gradient step
            self.optimizer.zero_grad()
            loss.mean().backward()
            self.optimizer.step()

        # Copy new weights into old policy
        self.policy_old.load_state_dict(self.policy.state_dict())

        # clear buffer
        self.buffer.clear()
        
        # Update weight
        self.diverse_weight_update_function()

    def save(self, checkpoint_path):
        torch.save(self.policy_old.state_dict(), checkpoint_path)

    def load(self, checkpoint_path):
        self.policy_old.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))
        self.policy.load_state_dict(torch.load(checkpoint_path, map_location=lambda storage, loc: storage))