<reponame>NLS-SJTU/BatcRL<gh_stars>1-10
import torch
import torch.nn as nn
import numpy as np
import os

class ReplayBuffer:
    def __init__(self, state_dim:int, max_size:int=10000, device=torch.device('cpu')):
        self.max_size = max_size
        self.device = device
        self.state_buffer = torch.empty((max_size, state_dim), dtype=torch.float32, device=device)
        # r_t, done,a_t, action_log
        self.other_buffer = torch.empty((max_size, 4), dtype=torch.float32, device=device)
        self.index = 0

    def append(self, state, other):
        self.index = self.index % self.max_size
        self.state_buffer[self.index] = torch.as_tensor(state, dtype=torch.float32, device=self.device)
        self.other_buffer[self.index] = torch.as_tensor(other, dtype=torch.float32, device=self.device)
        self.index += 1

    def sample_all(self, device):
        return (
            torch.as_tensor(self.state_buffer[:self.index], device=device), # s_t
            torch.as_tensor(self.other_buffer[:self.index, 2], dtype=torch.long, device=device), # a_t
            torch.as_tensor(self.other_buffer[:self.index, 0], device=device), # r_t
            torch.as_tensor(self.other_buffer[:self.index, 1], device=device), # done
            torch.as_tensor(self.other_buffer[:self.index, 3],device=device), # log_prob
        )

    def empty_buffer(self):
        self.index = 0


class ActorPPO(nn.Module):
    def __init__(self, state_dim:int, action_dim:int, mid_dim=512):
        super(ActorPPO, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, action_dim)
        )
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, state):
        return self.net(state)

    def get_action_and_logprob(self, state):
        prob = self.softmax(self.net(state))
        dist = torch.distributions.Categorical(prob)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        return action, log_prob

    def get_logprob_entropy(self, state, action):
        prob = self.softmax(self.net(state))
        dist = torch.distributions.Categorical(prob)
        logprob = dist.log_prob(action)
        entropy = dist.entropy().mean()
        return logprob, entropy


class CriticPPO(nn.Module):
    def __init__(self, state_dim:int, mid_dim:int=512):
        super(CriticPPO, self).__init__()
        self.state_value = nn.Sequential(
            nn.Linear(state_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, 1)
        )

    def forward(self, state):
        return self.state_value(state)


class PPODiscreteAgent:
    def __init__(self, state_dim:int, action_dim:int):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.mid_dim = 256
        self.actor_lr = 2e-4
        self.critic_lr = 1e-4
        self.entropy_coef= 0.02
        self.max_buffer_size= 100000
        self.batch_size = 32
        self.clip_epsilon = 0.2
        self.target_step = 2048
        self.repeat_time = 4
        self.reward_scale = 1.
        self.tau = 2 ** -9  # soft update.
        self.gamma = 0.98  # discount factor.

        self.net_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.buffer_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.replay_buffer = ReplayBuffer(self.state_dim, self.max_buffer_size, self.buffer_device)

        self.actor =  ActorPPO(self.state_dim,self.action_dim, self.mid_dim).to(self.net_device)
        self.critic = CriticPPO(self.state_dim, self.mid_dim).to(self.net_device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.critic_lr)
        self.last_state = None
        self.mse_loss = nn.MSELoss()

    def select_action(self, state):
        state = torch.as_tensor(state, dtype=torch.float32, device=self.net_device)
        action, logprob = self.actor.get_action_and_logprob(state)
        return action.detach().cpu().numpy(), logprob.detach().cpu().numpy()

    def explore_env(self, env):
        state = env.reset() if self.last_state is None else self.last_state
        for i in range(self.target_step):
            action, logprob = self.select_action(state)
            state_, reward, done, _ = env.step(int(action))
            others = (reward * self.reward_scale, (1 - done) * self.gamma, action, logprob)
            self.replay_buffer.append(state, others)
            state = env.reset() if done else state_
        self.last_state = state
        return self.target_step

    @staticmethod
    def optim_update(optimizer, objective):
        optimizer.zero_grad()
        objective.backward()
        optimizer.step()


    def update(self):
        with torch.no_grad():
            states, actions, reward, mask, old_logprob = self.replay_buffer.sample_all(self.net_device)
            state_values = self.critic(states)
            discount_cumulative_rewards = torch.empty(len(actions), dtype=torch.float32, device=self.net_device)
            tmp_last_state = torch.as_tensor((self.last_state,), dtype=torch.float32, device=self.net_device)
            last_value = self.critic(tmp_last_state)
            for i in range(len(actions)-1, -1, -1):
                discount_cumulative_rewards[i] = reward[i] + mask[i] * last_value
                last_value = discount_cumulative_rewards[i]

            advantage = discount_cumulative_rewards - (mask * state_values.squeeze(1))
            advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-10)
        for _ in range(int(self.target_step * self.repeat_time / self.batch_size)):
            indices = np.random.randint(0, len(actions), self.batch_size)
            batch_state = states[indices]
            batch_action = actions[indices]
            batch_dis_cum_rewards = discount_cumulative_rewards[indices]
            batch_advantage = advantage[indices]
            batch_old_logprob = old_logprob[indices]
            batch_new_logprob, batch_entropy = self.actor.get_logprob_entropy(batch_state, batch_action)
            ratio = (batch_new_logprob - batch_old_logprob).exp()
            surrogate1 = batch_advantage * ratio
            surrogate2 = batch_advantage * ratio.clamp(1 - self.clip_epsilon, 1 + self.clip_epsilon)
            obj_surrogate = -torch.min(surrogate1, surrogate2).mean()
            actor_loss = obj_surrogate + batch_entropy * self.entropy_coef
            self.optim_update(self.actor_optimizer, actor_loss)
            batch_values = self.critic(batch_state).squeeze(1)
            critic_loss = self.mse_loss(batch_values, batch_dis_cum_rewards)
            self.optim_update(self.critic_optimizer, critic_loss)
        self.replay_buffer.empty_buffer()


    def evaluate(self, env, render=False):
        epochs = 20
        res = np.zeros((epochs,))
        obs = env.reset()
        index = 0
        while index < epochs:
            if render: env.render()
            action, _ = self.select_action(obs)
            s_, reward, done, _ = env.step(int(action))
            res[index] += reward
            if done:
                index += 1
                obs = env.reset()
            else:
                obs = s_
        return res.mean(), res.std()

    def load_and_save_weight(self, path, mode='load'):
        actor_path = os.path.join(path, 'actor.pth')
        critic_path = os.path.join(path, 'critic.pth')
        if mode == 'load':
            if os.path.exists(actor_path) and os.path.exists(critic_path):
                self.actor.load_state_dict(torch.load(actor_path))
                self.critic.load_state_dict(torch.load(critic_path))
        else:
            if not os.path.exists(path):
                os.makedirs(path)
            torch.save(self.actor.state_dict(), actor_path)
            torch.save(self.critic.state_dict(), critic_path)















import torch
import torch.nn as nn
import numpy as np
import os

class ReplayBuffer:
    def __init__(self, state_dim:int, max_size:int=10000, device=torch.device('cpu')):
        self.max_size = max_size
        self.device = device
        self.state_buffer = torch.empty((max_size, state_dim), dtype=torch.float32, device=device)
        # r_t, done,a_t, action_log
        self.other_buffer = torch.empty((max_size, 4), dtype=torch.float32, device=device)
        self.index = 0

    def append(self, state, other):
        self.index = self.index % self.max_size
        self.state_buffer[self.index] = torch.as_tensor(state, dtype=torch.float32, device=self.device)
        self.other_buffer[self.index] = torch.as_tensor(other, dtype=torch.float32, device=self.device)
        self.index += 1

    def sample_all(self, device):
        return (
            torch.as_tensor(self.state_buffer[:self.index], device=device), # s_t
            torch.as_tensor(self.other_buffer[:self.index, 2], dtype=torch.long, device=device), # a_t
            torch.as_tensor(self.other_buffer[:self.index, 0], device=device), # r_t
            torch.as_tensor(self.other_buffer[:self.index, 1], device=device), # done
            torch.as_tensor(self.other_buffer[:self.index, 3],device=device), # log_prob
        )

    def empty_buffer(self):
        self.index = 0


class ActorPPO(nn.Module):
    def __init__(self, state_dim:int, action_dim:int, mid_dim=512):
        super(ActorPPO, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, action_dim)
        )
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, state):
        return self.net(state)

    def get_action_and_logprob(self, state):
        prob = self.softmax(self.net(state))
        dist = torch.distributions.Categorical(prob)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        return action, log_prob

    def get_logprob_entropy(self, state, action):
        prob = self.softmax(self.net(state))
        dist = torch.distributions.Categorical(prob)
        logprob = dist.log_prob(action)
        entropy = dist.entropy().mean()
        return logprob, entropy


class CriticPPO(nn.Module):
    def __init__(self, state_dim:int, mid_dim:int=512):
        super(CriticPPO, self).__init__()
        self.state_value = nn.Sequential(
            nn.Linear(state_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, mid_dim), nn.ReLU(),
            nn.Linear(mid_dim, 1)
        )

    def forward(self, state):
        return self.state_value(state)


class PPODiscreteAgent:
    def __init__(self, state_dim:int, action_dim:int):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.mid_dim = 256
        self.actor_lr = 2e-4
        self.critic_lr = 1e-4
        self.entropy_coef= 0.02
        self.max_buffer_size= 100000
        self.batch_size = 32
        self.clip_epsilon = 0.2
        self.target_step = 2048
        self.repeat_time = 4
        self.reward_scale = 1.
        self.tau = 2 ** -9  # soft update.
        self.gamma = 0.98  # discount factor.

        self.net_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.buffer_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.replay_buffer = ReplayBuffer(self.state_dim, self.max_buffer_size, self.buffer_device)

        self.actor =  ActorPPO(self.state_dim,self.action_dim, self.mid_dim).to(self.net_device)
        self.critic = CriticPPO(self.state_dim, self.mid_dim).to(self.net_device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=self.actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=self.critic_lr)
        self.last_state = None
        self.mse_loss = nn.MSELoss()

    def select_action(self, state):
        state = torch.as_tensor(state, dtype=torch.float32, device=self.net_device)
        action, logprob = self.actor.get_action_and_logprob(state)
        return action.detach().cpu().numpy(), logprob.detach().cpu().numpy()

    def explore_env(self, env):
        state = env.reset() if self.last_state is None else self.last_state
        for i in range(self.target_step):
            action, logprob = self.select_action(state)
            state_, reward, done, _ = env.step(int(action))
            others = (reward * self.reward_scale, (1 - done) * self.gamma, action, logprob)
            self.replay_buffer.append(state, others)
            state = env.reset() if done else state_
        self.last_state = state
        return self.target_step

    @staticmethod
    def optim_update(optimizer, objective):
        optimizer.zero_grad()
        objective.backward()
        optimizer.step()


    def update(self):
        with torch.no_grad():
            states, actions, reward, mask, old_logprob = self.replay_buffer.sample_all(self.net_device)
            state_values = self.critic(states)
            discount_cumulative_rewards = torch.empty(len(actions), dtype=torch.float32, device=self.net_device)
            tmp_last_state = torch.as_tensor((self.last_state,), dtype=torch.float32, device=self.net_device)
            last_value = self.critic(tmp_last_state)
            for i in range(len(actions)-1, -1, -1):
                discount_cumulative_rewards[i] = reward[i] + mask[i] * last_value
                last_value = discount_cumulative_rewards[i]

            advantage = discount_cumulative_rewards - (mask * state_values.squeeze(1))
            advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-10)
        for _ in range(int(self.target_step * self.repeat_time / self.batch_size)):
            indices = np.random.randint(0, len(actions), self.batch_size)
            batch_state = states[indices]
            batch_action = actions[indices]
            batch_dis_cum_rewards = discount_cumulative_rewards[indices]
            batch_advantage = advantage[indices]
            batch_old_logprob = old_logprob[indices]
            batch_new_logprob, batch_entropy = self.actor.get_logprob_entropy(batch_state, batch_action)
            ratio = (batch_new_logprob - batch_old_logprob).exp()
            surrogate1 = batch_advantage * ratio
            surrogate2 = batch_advantage * ratio.clamp(1 - self.clip_epsilon, 1 + self.clip_epsilon)
            obj_surrogate = -torch.min(surrogate1, surrogate2).mean()
            actor_loss = obj_surrogate + batch_entropy * self.entropy_coef
            self.optim_update(self.actor_optimizer, actor_loss)
            batch_values = self.critic(batch_state).squeeze(1)
            critic_loss = self.mse_loss(batch_values, batch_dis_cum_rewards)
            self.optim_update(self.critic_optimizer, critic_loss)
        self.replay_buffer.empty_buffer()


    def evaluate(self, env, render=False):
        epochs = 20
        res = np.zeros((epochs,))
        obs = env.reset()
        index = 0
        while index < epochs:
            if render: env.render()
            action, _ = self.select_action(obs)
            s_, reward, done, _ = env.step(int(action))
            res[index] += reward
            if done:
                index += 1
                obs = env.reset()
            else:
                obs = s_
        return res.mean(), res.std()

    def load_and_save_weight(self, path, mode='load'):
        actor_path = os.path.join(path, 'actor.pth')
        critic_path = os.path.join(path, 'critic.pth')
        if mode == 'load':
            if os.path.exists(actor_path) and os.path.exists(critic_path):
                self.actor.load_state_dict(torch.load(actor_path))
                self.critic.load_state_dict(torch.load(critic_path))
        else:
            if not os.path.exists(path):
                os.makedirs(path)
            torch.save(self.actor.state_dict(), actor_path)
            torch.save(self.critic.state_dict(), critic_path)
# r_t, done,a_t, action_log # s_t # a_t # r_t # done # log_prob # soft update. # discount factor.