04 - 策略梯度方法¶

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

class PolicyNetwork(nn.Module):  # 继承nn.Module定义网络层
    """策略网络"""

    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super(PolicyNetwork, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Softmax(dim=-1)
        )

    def forward(self, state):
        """返回动作概率分布"""
        return self.net(state)

    def select_action(self, state):
        """采样动作"""
        probs = self.forward(state)
        # 数值稳定性: 加小 epsilon 防止 log(0) 产生 -inf
        probs = probs + 1e-8
        dist = Categorical(probs)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        return action.item(), log_prob  # 将单元素张量转为Python数值

class REINFORCE:
    """REINFORCE算法"""

    def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99):
        self.gamma = gamma

        # 策略网络
        self.policy = PolicyNetwork(state_dim, action_dim)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)

    def compute_returns(self, rewards):
        """计算折扣回报"""
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + self.gamma * G
            returns.insert(0, G)

        # 归一化回报（减少方差）
        returns = torch.tensor(returns, dtype=torch.float32)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        return returns

    def update(self, log_probs, rewards):
        """策略更新"""
        returns = self.compute_returns(rewards)

        # 计算策略损失
        policy_loss = []
        for log_prob, G in zip(log_probs, returns):  # zip按位置配对
            policy_loss.append(-log_prob * G)  # 负号因为梯度上升

        loss = torch.stack(policy_loss).sum()  # torch.stack沿新维度拼接张量

        # 优化
        self.optimizer.zero_grad()  # 清零梯度
        loss.backward()  # 反向传播计算梯度
        self.optimizer.step()  # 更新参数

        return loss.item()

    def train(self, env, num_episodes=1000):
        """训练"""
        rewards_history = []

        for episode in range(num_episodes):
            state, _ = env.reset()
            log_probs = []
            rewards = []
            done = False

            # 收集轨迹
            while not done:
                state_tensor = torch.FloatTensor(state)
                action, log_prob = self.policy.select_action(state_tensor)

                next_state, reward, terminated, truncated, _ = env.step(action)
                done = terminated or truncated

                log_probs.append(log_prob)
                rewards.append(reward)
                state = next_state

            # 更新策略
            loss = self.update(log_probs, rewards)

            total_reward = sum(rewards)
            rewards_history.append(total_reward)

            if (episode + 1) % 100 == 0:
                avg_reward = np.mean(rewards_history[-100:])
                print(f"Episode {episode + 1}, Avg Reward: {avg_reward:.2f}")

        return rewards_history

class ValueNetwork(nn.Module):
    """值函数网络"""

    def __init__(self, state_dim, hidden_dim=128):
        super(ValueNetwork, self).__init__()

        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )

    def forward(self, state):
        return self.net(state).squeeze(-1)  # squeeze压缩维度

class REINFORCEWithBaseline:
    """带基线的REINFORCE"""

    def __init__(self, state_dim, action_dim, lr_policy=1e-3,
                 lr_value=1e-3, gamma=0.99):
        self.gamma = gamma

        # 策略网络
        self.policy = PolicyNetwork(state_dim, action_dim)
        self.policy_optimizer = optim.Adam(self.policy.parameters(), lr=lr_policy)

        # 值网络（基线）
        self.value = ValueNetwork(state_dim)
        self.value_optimizer = optim.Adam(self.value.parameters(), lr=lr_value)

    def update(self, states, log_probs, rewards):
        """更新策略和值函数"""
        states = torch.stack([torch.FloatTensor(s) for s in states])

        # 计算回报
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + self.gamma * G
            returns.insert(0, G)
        returns = torch.tensor(returns, dtype=torch.float32)

        # 更新值函数
        values = self.value(states)
        value_loss = nn.MSELoss()(values, returns)

        self.value_optimizer.zero_grad()
        value_loss.backward()
        self.value_optimizer.step()

        # 更新策略（使用优势函数）
        advantages = returns - values.detach()  # 分离计算图，不参与梯度计算

        policy_loss = []
        for log_prob, advantage in zip(log_probs, advantages):
            policy_loss.append(-log_prob * advantage)

        policy_loss = torch.stack(policy_loss).sum()

        self.policy_optimizer.zero_grad()
        policy_loss.backward()
        self.policy_optimizer.step()

        return policy_loss.item(), value_loss.item()

class ActorCritic(nn.Module):
    """Actor-Critic网络"""

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

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

        # Actor头
        self.actor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Softmax(dim=-1)
        )

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

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

    def select_action(self, state):
        """选择动作"""
        probs, value = self.forward(state)
        dist = Categorical(probs)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        return action.item(), log_prob, value

class A2CAgent:
    """A2C智能体"""

    def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99,
                 value_coef=0.5, entropy_coef=0.01):
        self.gamma = gamma
        self.value_coef = value_coef
        self.entropy_coef = entropy_coef

        self.model = ActorCritic(state_dim, action_dim)
        self.optimizer = optim.Adam(self.model.parameters(), lr=lr)

    def compute_returns(self, rewards, values, done):
        """计算回报"""
        returns = []
        G = 0 if done else values[-1].item()  # [-1]负索引取最后元素

        for r, v in zip(reversed(rewards), reversed(values[:-1])):
            G = r + self.gamma * G
            returns.insert(0, G)

        return torch.tensor(returns, dtype=torch.float32)

    def update(self, states, actions, log_probs, rewards, values, done):
        """更新Actor和Critic"""
        states = torch.stack(states)
        actions = torch.tensor(actions)
        old_log_probs = torch.stack(log_probs)
        values = torch.stack(values).squeeze(-1)

        # 计算回报
        returns = self.compute_returns(rewards, values, done)

        # 重新计算（用于获取新的log_prob和entropy）
        probs, new_values = self.model(states)
        dist = Categorical(probs)
        new_log_probs = dist.log_prob(actions)
        entropy = dist.entropy()

        # 优势函数
        advantages = returns - new_values.squeeze(-1).detach()

        # Actor损失（策略梯度）
        actor_loss = -(new_log_probs * advantages).mean()

        # Critic损失（值函数）
        critic_loss = nn.MSELoss()(new_values.squeeze(-1), returns)

        # 熵奖励（鼓励探索）
        entropy_loss = -entropy.mean()

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

        # 优化
        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)
        self.optimizer.step()

        return loss.item()