text: from agents.agent import Agent
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from torch.distributions import Categorical
  from torch.utils.data import SubsetRandomSampler, BatchSampler
  import numpy as np
  
  
  class PPOAgent(Agent):
      def __init__(self, state_space: int, action_space: int, a_hidden=50, c_hidden=50,
                   a_lr=1e-3, c_lr=3e-3, gamma=0.9, clip_e=0.2):
          self.state_space = state_space
          self.action_space = action_space
          self.gamma = gamma
          self.clip_e = clip_e
  
          self.actor = Actor(state_space, a_hidden, action_space)
          self.actor_optim = optim.Adam(self.actor.parameters(), lr=a_lr)
  
          self.critic = Critic(state_space, c_hidden)
          self.critic_optim = optim.Adam(self.critic.parameters(), lr=c_lr)
  
          self.states = []
          self.new_states = []
          self.actions = []
          self.a_probs = []
          self.rewards = []
  
      def act(self, state):
          state = torch.tensor(state, dtype=torch.float).view(-1, self.state_space)  # 1st dimension is batch number
          with torch.no_grad():
              probs = self.actor(state)
          # print(probs)
          dist = Categorical(probs)
          action = dist.sample()
          return action.item(), probs[:, action.item()].item()
  
      def train(self, batch_size=8):
          # PPO algorithm
          # Unroll rewards
          rewards = np.array(self.rewards)
          reward = 0
          for i in reversed(range(len(self.rewards))):
              rewards[i] += self.gamma * reward
              reward = rewards[i]
  
          states = torch.tensor(self.states, dtype=torch.float)
          actions = torch.tensor(self.actions, dtype=torch.long).view(-1, 1)
          old_probs = torch.tensor(self.a_probs, dtype=torch.float).view(-1, 1)
          rewards = torch.tensor(rewards, dtype=torch.float).view(-1, 1)
  
          for batch in BatchSampler(SubsetRandomSampler(range(len(self.states))), batch_size, drop_last=False):
              # Calculate advantage
              Rt = rewards[batch]
              V = self.critic(states[batch])
              advantage = Rt - V
              advantage.detach_()     # Inplace detach autograd
  
              # Calculate PPO loss
              new_probs = self.actor(states[batch]).gather(1, actions[batch])
              prob_ratio = new_probs / old_probs[batch]
              surr = prob_ratio * advantage
              clipped_surr = torch.clamp(prob_ratio, 1 - self.clip_e, 1 + self.clip_e) * advantage
              actor_loss = -torch.min(surr, clipped_surr).mean()  # Mean is here for when batch_size > 1
  
              # Update actor
              self.actor_optim.zero_grad()
              actor_loss.backward()
              self.actor_optim.step()
  
              # Update critic
              critic_loss = F.mse_loss(Rt, V)
              self.critic_optim.zero_grad()
              critic_loss.backward()
              self.critic_optim.step()
  
          self._clear_buffers()
  
      def store_transition(self, state, new_state, action, a_prob, reward):
          self.states.append(state)
          self.new_states.append(new_state)
          self.actions.append(action)
          self.a_probs.append(a_prob)
          self.rewards.append(reward)
  
      def _clear_buffers(self):
          self.states = []
          self.new_states = []
          self.actions = []
          self.a_probs = []
          self.rewards = []
  
  
  class Actor(nn.Module):
      def __init__(self, input, hidden, output):
          super(Actor, self).__init__()
  
          self.fc1 = nn.Linear(input, hidden)
          self.fc2 = nn.Linear(hidden, output)
  
      def forward(self, x):
          x = F.relu(self.fc1(x))
          x = F.softmax(self.fc2(x), dim=1)
          return x
  
  
  class Critic(nn.Module):
      def __init__(self, input, hidden):
          super(Critic, self).__init__()
  
          self.fc1 = nn.Linear(input, hidden)
          self.fc2 = nn.Linear(hidden, 1)
  
      def forward(self, x):
          x = F.relu(self.fc1(x))
          x = self.fc2(x)
          return x
  