pytorch ppo代码示例

PyTorch PPO（Proximal Policy Optimization）是一种常用的强化学习算法，用于训练策略模型。下面是一个简单的PyTorch PPO代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
import gym

# 定义策略网络
class Policy(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Policy, self).__init__()
        self.fc = nn.Linear(input_dim, 64)
        self.actor = nn.Linear(64, output_dim)
        self.critic = nn.Linear(64, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc(x))
        action_probs = torch.softmax(self.actor(x), dim=-1)
        state_value = self.critic(x)
        return action_probs, state_value

# 定义PPO算法
class PPO:
    def __init__(self, input_dim, output_dim):
        self.policy = Policy(input_dim, output_dim)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=0.001)
    
    def select_action(self, state):
        state = torch.FloatTensor(state)
        action_probs, _ = self.policy(state)
        action_dist = torch.distributions.Categorical(action_probs)
        action = action_dist.sample()
        return action.item()
    
    def update(self, states, actions, log_probs, returns, advantages):
        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions)
        log_probs = torch.FloatTensor(log_probs)
        returns = torch.FloatTensor(returns)
        advantages = torch.FloatTensor(advantages)
        
        # 计算策略损失和价值损失
        action_probs, state_values = self.policy(states)
        dist = torch.distributions.Categorical(action_probs)
        new_log_probs = dist.log_prob(actions)
        ratio = torch.exp(new_log_probs - log_probs)
        surr1 = ratio * advantages
        surr2 = torch.clamp(ratio, 1-0.2, 1+0.2) * advantages
        policy_loss = -torch.min(surr1, surr2).mean()
        value_loss = nn.MSELoss()(state_values, returns.unsqueeze(1))
        
        # 更新策略网络
        self.optimizer.zero_grad()
        loss = policy_loss + 0.5 * value_loss
        loss.backward()
        self.optimizer.step()

# 创建环境和PPO对象
env = gym.make('CartPole-v1')
input_dim = env.observation_space.shape
output_dim = env.action_space.n
ppo = PPO(input_dim, output_dim)

# 训练PPO模型
max_episodes = 1000
max_steps = 200
for episode in range(max_episodes):
    state = env.reset()
    states, actions, log_probs, rewards = [], [], [], []
    for step in range(max_steps):
        action = ppo.select_action(state)
        next_state, reward, done, _ = env.step(action)
        
        states.append(state)
        actions.append(action)
        log_probs.append(torch.log(ppo.policy(torch.FloatTensor(state))[action]))
        rewards.append(reward)
        
        state = next_state
        
        if done:
            break
    
    # 计算回报和优势值
    returns = []
    advantages = []
    G = 0
    for r in reversed(rewards):
        G = r + 0.99 * G
        returns.insert(0, G)
    returns = torch.tensor(returns)
    returns = (returns -9)
    
    for t in range(len(rewards)):
        advantages.append(returns[t] - ppo.policy(torch.FloatTensor(states[t])))
    advantages = torch.tensor(advantages)
    
    # 更新策略网络
    ppo.update(states, actions, log_probs, returns, advantages)

# 使用训练好的模型进行测试
state = env.reset()
total_reward = 0
while True:
    env.render()
    action = ppo.select_action(state)
    state, reward, done, _ = env.step(action)
    total_reward += reward
    if done:
        break

print("Total reward:", total_reward)

这个示例代码使用PyTorch实现了一个简单的PPO算法，用于在CartPole-v1环境中训练一个策略模型。代码中包含了策略网络的定义、PPO算法的实现以及训练和测试的过程。