assistant: ```
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.softmax(self.fc2(x), dim=-1)
        return x


def ppo_loss(ratio, advantage, clip_ratio, entropy, c1, c2):
    pg_loss = torch.min(ratio * advantage, torch.clamp(ratio, 1-clip_ratio, 1+clip_ratio) * advantage)
    vf_loss = torch.square(ratio - 1.0) * advantage
    entropy_loss = -c2 * entropy
    loss = -pg_loss + c1 * vf_loss + entropy_loss
    return loss.mean()


class PPO:
    def __init__(self, state_size, action_size, hidden_size=64, clip_ratio=0.2, c1=1.0, c2=0.01, gamma=0.99, lam=0.95, lr=3e-4):
        self.actor = MLP(state_size, hidden_size, action_size)
        self.critic = MLP(state_size, hidden_size, 1)
        self.optimizer = optim.Adam(list(self.actor.parameters())+list(self.critic.parameters()), lr=lr)
        self.clip_ratio = clip_ratio
        self.c1 = c1
        self.c2 = c2
        self.gamma = gamma
        self.lam = lam
        
    def select_action(self, state):
        state = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
        action_probs = self.actor(state)
        dist = torch.distributions.Categorical(action_probs)
        action = dist.sample()
        return action.item(), dist.log_prob(action), dist.entropy()
    
    def evaluate(self, states, actions):
        action_probs = self.actor(states)
        dist = torch.distributions.Categorical(action_probs)
        log_probs = dist.log_prob(actions)
        entropy = dist.entropy()
        values = self.critic(states)
        return log_probs, entropy, values
        
    def train(self, states, actions, rewards, next_states, masks):
        states = torch.tensor(states, dtype=torch.float32)
        actions = torch.tensor(actions, dtype=torch.int32)
        rewards = torch.tensor(rewards, dtype=torch.float32)

        _, _, values = self.evaluate(states, actions)
        _, _, next_values = self.evaluate(next_states, actions)
        deltas = rewards + self.gamma * masks * next_values - values
        advantages = compute_advantage(deltas, masks, self.gamma, self.lam)

        for _ in range(num_epochs):
            indices = torch.randperm(states.size(0))
            for start in range(0, states.size(0), batch_size):
                end = min(start + batch_size, states.size(0))
                mini_batch_indices = indices[start:end]
                mini_batch_states = states[mini_batch_indices]
                mini_batch_actions = actions[mini_batch_indices]
                mini_batch_advantages = advantages[mini_batch_indices]
                mini_batch_old_log_probs, mini_batch_entropy, _ = self.evaluate(mini_batch_states, mini_batch_actions)
                for _ in range(num_mini_batches):
                    mini_batch_indices = torch.randperm(mini_batch_states.size(0))
                    for start in range(0, mini_batch_states.size(0), mini_batch_size):
                        end = min(start + mini_batch_size, mini_batch_states.size(0))
                        mini_batch_indices_ = mini_batch_indices[start:end]
                        mini_batch_states_ = mini_batch_states[mini_batch_indices_]
                        mini_batch_actions_ = mini_batch_actions[mini_batch_indices_]
                        mini_batch_advantages_ = mini_batch_advantages[mini_batch_indices_]
                        mini_batch_old_log_probs_, mini_batch_entropy_, mini_batch_values_ = self.evaluate(mini_batch_states_, mini_batch_actions_)
                        ratio = torch.exp(mini_batch_old_log_probs_ - mini_batch_old_log_probs.detach())
                        loss = ppo_loss(ratio, mini_batch_advantages_, self.clip_ratio, mini_batch_entropy_, self.c1, self.c2)
                        self.optimizer.zero_grad()
                        loss.backward()
                        self.optimizer.step()


def compute_advantage(deltas, masks, gamma, lam):
    advantages = torch.zeros_like(deltas)
    last_advantage = 0
    for t in reversed(range(deltas.size(0))):
        if t == deltas.size(0) - 1:
            next_non_terminal = 1.0 - masks[-1]
            next_value = last_advantage = deltas[-1] + next_non_terminal * last_advantage * gamma
        else:
            next_non_terminal = 1.0 - masks[t]
            next_value = last_advantage = deltas[t] + next_non_terminal * last_advantage * gamma * lam
        advantages[t] = next_value - self.critic(states[t])
    return advantages
```