05 - Actor-Critic高级方法¶

def compute_gae(rewards, values, dones, gamma=0.99, lambda_=0.95, next_value=0):
    """
    计算GAE

    参数:
        rewards: 奖励序列（长度N）
        values: 值函数估计序列（长度N，对应每个时间步的V(s_t)）
        dones: 终止标志序列（长度N）
        gamma: 折扣因子
        lambda_: GAE参数
        next_value: 最后一步之后的bootstrap值V(s_{N})，终止时为0

    返回:
        advantages: 优势估计
        returns: 回报估计
    """
    advantages = []
    gae = 0

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

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

    advantages = np.array(advantages)  # np.array创建NumPy数组
    returns = advantages + np.array(values)

    return advantages, returns

import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from torch.distributions import Categorical

class A3CNetwork(nn.Module):  # 继承nn.Module定义网络层
    """A3C共享网络"""

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

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

        # 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

class A3CWorker(mp.Process):
    """A3C Worker进程"""

    def __init__(self, global_net, optimizer, env_name, worker_id,
                 gamma=0.99, max_steps=20):
        super(A3CWorker, self).__init__()

        self.global_net = global_net
        self.optimizer = optimizer
        self.env_name = env_name
        self.worker_id = worker_id
        self.gamma = gamma
        self.max_steps = max_steps

        # 本地网络
        self.local_net = A3CNetwork(global_net.shared[0].in_features,
                                     global_net.actor.out_features)
        self.local_net.load_state_dict(global_net.state_dict())

    def run(self):
        """Worker主循环"""
        env = gym.make(self.env_name)
        state, _ = env.reset()

        while True:
            # 同步全局参数
            self.local_net.load_state_dict(self.global_net.state_dict())

            # 收集经验
            states, actions, rewards, values, log_probs = [], [], [], [], []

            for step in range(self.max_steps):
                state_tensor = torch.FloatTensor(state)
                logits, value = self.local_net(state_tensor)

                dist = Categorical(logits=logits)
                action = dist.sample()
                log_prob = dist.log_prob(action)

                next_state, reward, terminated, truncated, _ = env.step(action.item())  # 将单元素张量转为Python数值
                done = terminated or truncated

                states.append(state)
                actions.append(action)
                rewards.append(reward)
                values.append(value)
                log_probs.append(log_prob)

                state = next_state

                if done:
                    state, _ = env.reset()
                    break

            # 计算回报和优势
            R = 0 if done else self.local_net(
                torch.FloatTensor(state))[1].item()

            returns = []
            for r in reversed(rewards):
                R = r + self.gamma * R
                returns.insert(0, R)

            returns = torch.tensor(returns)
            values = torch.stack(values).squeeze()  # squeeze压缩维度  # torch.stack沿新维度拼接张量
            log_probs = torch.stack(log_probs)

            advantages = returns - values.detach()  # 分离计算图，不参与梯度计算

            # 计算损失
            actor_loss = -(log_probs * advantages).mean()
            critic_loss = nn.MSELoss()(values, returns)
            entropy = Categorical(logits=logits).entropy().mean()

            loss = actor_loss + 0.5 * critic_loss - 0.01 * entropy

            # 异步更新全局网络
            self.optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度

            # 梯度上传
            for local_param, global_param in zip(  # zip按位置配对
                self.local_net.parameters(),
                self.global_net.parameters()
            ):
                if global_param.grad is not None:
                    global_param.grad += local_param.grad
                else:
                    global_param.grad = local_param.grad.clone()

            self.optimizer.step()  # 更新参数

class A3C:
    """A3C主类"""

    def __init__(self, state_dim, action_dim, lr=1e-4, num_workers=4):
        self.global_net = A3CNetwork(state_dim, action_dim)
        self.global_net.share_memory()

        self.optimizer = optim.Adam(self.global_net.parameters(), lr=lr)
        self.num_workers = num_workers

    def train(self, env_name, num_episodes=10000):
        """训练"""
        workers = [
            A3CWorker(self.global_net, self.optimizer, env_name, i)
            for i in range(self.num_workers)
        ]

        for worker in workers:
            worker.start()

        for worker in workers:
            worker.join()