01 - PPO算法：近端策略优化¶

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.distributions import Categorical

class ActorCritic(nn.Module):  # 继承nn.Module定义网络层
    """Actor-Critic网络"""

    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super(ActorCritic, self).__init__()

        # 共享层
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh()
        )

        # Actor头
        self.actor = nn.Linear(hidden_dim, action_dim)

        # Critic头
        self.critic = nn.Linear(hidden_dim, 1)

    def forward(self, state):
        features = self.shared(state)
        logits = self.actor(features)
        value = self.critic(features)
        return logits, value

    def get_action_and_value(self, state, action=None):
        """获取动作、对数概率和价值"""
        logits, value = self.forward(state)
        dist = Categorical(logits=logits)

        if action is None:
            action = dist.sample()

        log_prob = dist.log_prob(action)
        entropy = dist.entropy()

        return action, log_prob, entropy, value.squeeze(-1)  # squeeze压缩维度

class PPO:
    """PPO算法"""

    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99,
                 gae_lambda=0.95, clip_epsilon=0.2,
                 value_coef=0.5, entropy_coef=0.01,
                 max_grad_norm=0.5, update_epochs=10, batch_size=64):

        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.clip_epsilon = clip_epsilon
        self.value_coef = value_coef
        self.entropy_coef = entropy_coef
        self.max_grad_norm = max_grad_norm
        self.update_epochs = update_epochs
        self.batch_size = batch_size

        # 网络
        self.ac = ActorCritic(state_dim, action_dim)
        self.optimizer = optim.Adam(self.ac.parameters(), lr=lr)

    def compute_gae(self, rewards, values, dones, next_value):
        """计算广义优势估计（GAE）"""
        advantages = []
        gae = 0

        for t in reversed(range(len(rewards))):
            if t == len(rewards) - 1:
                next_val = next_value
            else:
                next_val = values[t + 1]

            delta = rewards[t] + self.gamma * next_val * (1 - dones[t]) - values[t]
            gae = delta + self.gamma * self.gae_lambda * (1 - dones[t]) * gae
            advantages.insert(0, gae)

        advantages = torch.tensor(advantages, dtype=torch.float32)
        returns = advantages + torch.tensor(values, dtype=torch.float32)

        return advantages, returns

    def update(self, states, actions, old_log_probs, rewards, values, dones, next_value):
        """PPO更新"""
        # 计算优势函数和回报
        advantages, returns = self.compute_gae(rewards, values, dones, next_value)

        # 归一化优势
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # 转换为张量
        states = torch.stack(states)  # torch.stack沿新维度拼接张量
        actions = torch.tensor(actions)
        old_log_probs = torch.stack(old_log_probs)

        # 多轮更新
        for _ in range(self.update_epochs):
            # 重新计算log_prob和价值
            _, new_log_probs, entropy, new_values = self.ac.get_action_and_value(states, actions)

            # 概率比率
            ratio = torch.exp(new_log_probs - old_log_probs)

            # 裁剪目标
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages
            actor_loss = -torch.min(surr1, surr2).mean()

            # 值函数损失
            value_loss = nn.MSELoss()(new_values, returns)

            # 熵奖励
            entropy_loss = -entropy.mean()

            # 总损失
            loss = actor_loss + self.value_coef * value_loss + self.entropy_coef * entropy_loss

            # 优化
            self.optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            nn.utils.clip_grad_norm_(self.ac.parameters(), self.max_grad_norm)
            self.optimizer.step()  # 更新参数

        return loss.item()  # 将单元素张量转为Python数值

    def select_action(self, state):
        """选择动作"""
        with torch.no_grad():  # 禁用梯度计算，节省内存
            state = torch.FloatTensor(state)
            action, log_prob, _, value = self.ac.get_action_and_value(state)
            return action.item(), log_prob, value.item()