ppo算法代码pytorch

下面是一个使用PyTorch实现的PPO算法的代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(ActorCritic, self).__init__()
        self.fc1 = nn.Linear(state_dim, 128)
        self.fc2 = nn.Linear(128, 64)
        
        self.actor = nn.Linear(64, action_dim)
        self.critic = nn.Linear(64, 1)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        
        action_probs = F.softmax(self.actor(x), dim=-1)
        state_value = self.critic(x)
        
        return action_probs, state_value

class PPO:
    def __init__(self, state_dim, action_dim, lr_actor=0.0003, lr_critic=0.001, gamma=0.99, clip_ratio=0.2):
        self.actor_critic = ActorCritic(state_dim, action_dim)
        self.optimizer_actor = optim.Adam(self.actor_critic.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.actor_critic.critic.parameters(), lr=lr_critic)
        self.gamma = gamma
        self.clip_ratio = clip_ratio
        
    def compute_returns(self, rewards, masks, values):
        returns = torch.zeros_like(rewards)
        discounted_return = 0
        for i in reversed(range(len(rewards))):
            discounted_return = rewards[i] + self.gamma * discounted_return * masks[i]
            returns[i] = discounted_return
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        return returns
        
    def compute_advantage(self, rewards, masks, values):
        returns = self.compute_returns(rewards, masks, values).detach()
        advantages = returns - values
        return advantages
        
    def update(self, states, actions, old_log_probs, rewards, masks):
        action_probs, values = self.actor_critic(states)
        
        returns = self.compute_returns(rewards, masks, values)
        advantages = self.compute_advantage(rewards, masks, values)
        
        ratio = torch.exp(action_probs.log_prob(actions) - old_log_probs)
        surrogate1 = ratio * advantages
        surrogate2 = torch.clamp(ratio, 1 - self.clip_ratio, 1 + self.clip_ratio) * advantages
        
        actor_loss = -torch.min(surrogate1, surrogate2).mean()
        critic_loss = F.mse_loss(returns, values)
        
        self.optimizer_actor.zero_grad()
        actor_loss.backward()
        self.optimizer_actor.step()
        
        self.optimizer_critic.zero_grad()
        critic_loss.backward()
        self.optimizer_critic.step()

这里的代码实现了一个简单的ActorCritic模型作为PPO算法的基础。在`PPO`类中，`update`方法用于更新模型参数，`compute_returns`方法计算回报值，`compute_advantage`方法计算优势值。PPO算法的核心在于使用两个surrogate loss来进行策略优化。