PyTorch实战：用A2C算法训练CartPole游戏AI（附完整代码解析）

从零构建A2C智能体：在CartPole中掌握策略梯度与价值评估的协同艺术

如果你已经尝试过用传统Q-learning玩过一些简单的游戏环境，可能会发现当状态空间稍微复杂一点时，Q-table就会变得难以管理。而当你转向深度强化学习时，DDPG、PPO这些名词又让人望而生畏。其实，在深度策略梯度方法中，有一个既优雅又实用的算法，它巧妙地平衡了学习效率和实现复杂度——这就是Advantage Actor-Critic，也就是我们常说的A2C。

今天，我们不谈那些晦涩的数学推导，而是直接动手，用PyTorch从零开始构建一个完整的A2C智能体，让它学会玩经典的CartPole平衡游戏。我会带你走过每一个关键步骤：从环境搭建、网络设计，到训练循环中的每一个细节，最后还会分享几个我实际调试中发现的、能让训练更稳定的小技巧。无论你是想快速上手一个可运行的强化学习项目，还是希望深入理解Actor-Critic框架的实际运作，这篇文章都能给你带来实实在在的收获。

1. 环境配置与项目初始化

在开始写任何代码之前，我们需要先搭建好实验环境。CartPole是OpenAI Gym中最经典的测试环境之一，它的目标很简单：控制一个小车左右移动，让连接在小车上的杆子保持竖直不倒。这个环境状态空间小（4个维度），动作空间离散（2个动作），非常适合作为强化学习算法的“Hello World”。

1.1 安装必要的依赖库

我建议使用conda或venv创建一个独立的Python环境，避免包版本冲突。以下是核心依赖：

# 创建并激活虚拟环境（以conda为例）
conda create -n rl_a2c python=3.8
conda activate rl_a2c

# 安装核心库
pip install torch==1.9.0
pip install gym==0.21.0
pip install numpy==1.21.0
pip install matplotlib==3.4.0

这里我固定了几个主要库的版本，因为不同版本的gym或torch可能在API上有细微差别，固定版本能确保代码可复现。如果你用的是没有GPU的机器，安装PyTorch时可以使用pip install torch==1.9.0+cpu。

1.2 理解CartPole环境接口

在写智能体之前，我们先花几分钟熟悉一下gym环境的基本用法：

import gym

# 创建环境
env = gym.make('CartPole-v1')

# 重置环境，获取初始状态
state = env.reset()
print(f"状态空间维度: {env.observation_space.shape}")
print(f"动作空间大小: {env.action_space.n}")

# 执行一个随机动作
action = env.action_space.sample()
next_state, reward, done, info = env.step(action)
print(f"执行动作 {action} 后，奖励: {reward}, 是否结束: {done}")

运行这段代码，你会看到状态是一个4维向量，分别表示小车位置、小车速度、杆子角度和杆子顶端速度。动作只有两个：0（向左推）和1（向右推）。每步奖励都是+1，除非杆子倒下或小车超出边界，此时done变为True。

注意：CartPole-v1的最大步数是500步，达到500步后也会自动结束。这意味着完美策略最多能获得500分。

2. A2C算法核心思想解析

在深入代码之前，我们需要搞清楚A2C到底在做什么。很多人一听到“Actor-Critic”就觉得复杂，其实它的思想非常直观。

2.1 从策略梯度到优势函数

传统的REINFORCE策略梯度算法直接使用整个回合的总回报来更新策略，这就像你玩完一局游戏后，根据最终得分来调整每一步的决策。这种方法有两个明显问题：

1. 高方差：单次回合的得分波动很大，导致训练不稳定
2. 延迟奖励：只有回合结束后才能更新，学习效率低

Actor-Critic框架引入了一个“评论家”（Critic）来实时评估状态的价值，用这个评估值作为基准来调整“演员”（Actor）的策略。A2C中的“A”代表Advantage（优势），它衡量的是某个动作相对于平均水平的优势程度。

用公式表示就是：

优势值 = Q(s,a) - V(s)

其中Q(s,a)是执行动作a后的预期回报，V(s)是状态s的平均价值。如果优势值为正，说明这个动作比平均水平好，应该加强；如果为负，则应该减弱。

2.2 A2C与A3C的区别

你可能也听说过A3C（Asynchronous Advantage Actor-Critic）。A2C其实就是A3C的同步版本：

A2C通过同步更新避免了策略不一致的问题，在实际应用中往往更稳定，特别是在有GPU加速的情况下，批量处理数据效率更高。

2.3 算法流程概览

A2C的训练遵循一个清晰的循环：

1. 并行采样：多个环境实例同时运行，收集经验数据
2. 优势计算：用Critic网络估计状态价值，计算优势值
3. 策略更新：用优势值调整Actor网络（策略梯度）
4. 价值更新：更新Critic网络以更准确估计价值
5. 重复循环：直到策略收敛

这个流程的关键在于，我们同时学习两个东西：什么动作好（Actor）和状态有多好（Critic）。两者相互促进，就像教练和运动员的关系——运动员（Actor）尝试动作，教练（Critic）给出反馈，两者共同进步。

3. 网络架构设计与实现

现在我们来设计神经网络部分。A2C需要两个输出头：一个用于策略（Actor），一个用于价值估计（Critic）。但注意，我们通常共享底层的特征提取层。

3.1 构建Actor-Critic网络

我设计了一个简洁但有效的网络结构：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical

class ActorCritic(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim=256):
        super(ActorCritic, self).__init__()
        
        # 共享的特征提取层
        self.shared = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        
        # Actor分支：输出动作概率分布
        self.actor = nn.Sequential(
            nn.Linear(hidden_dim, output_dim),
            nn.Softmax(dim=-1)  # 转换为概率
        )
        
        # Critic分支：输出状态价值（单个标量）
        self.critic = nn.Linear(hidden_dim, 1)
    
    def forward(self, x):
        # 提取共享特征
        features = self.shared(x)
        
        # Actor输出：动作概率
        action_probs = self.actor(features)
        
        # Critic输出：状态价值
        state_value = self.critic(features)
        
        # 创建概率分布（方便采样和计算log概率）
        dist = Categorical(action_probs)
        
        return dist, state_value

这个设计有几个值得注意的地方：

• 共享层：Actor和Critic共享前几层，这不仅能减少参数量，还能让两者学习到相同的状态表示
• Softmax输出：Actor输出的是每个动作的概率，而不是直接输出动作。这符合策略梯度方法的随机策略特性
• Categorical分布：我们使用PyTorch的分布API，它能方便地处理采样、log概率计算等操作

3.2 为什么使用共享层？

有些实现会为Actor和Critic分别设计独立的网络，但我更推荐共享层设计，原因如下：

1. 计算效率：只需一次前向传播就能得到策略和价值
2. 特征一致性：Actor和Critic基于相同的状态表示进行决策和评估
3. 稳定训练：共享梯度有助于防止某个分支过拟合

不过要注意，如果Actor和Critic的任务差异太大（比如在复杂环境中），独立网络可能更好。对于CartPole这样的简单环境，共享层完全够用。

3.3 网络初始化技巧

神经网络的初始化对训练稳定性影响很大。我通常会在网络定义后添加自定义初始化：

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.orthogonal_(m.weight, gain=0.01)
        nn.init.constant_(m.bias, 0)

model = ActorCritic(input_dim=4, output_dim=2)
model.apply(init_weights)

这里使用正交初始化（orthogonal initialization）而不是常见的Xavier或He初始化，因为在强化学习中，正交初始化往往能带来更稳定的训练。gain参数设为较小的值（如0.01）可以防止初始输出过大。

4. 训练循环的完整实现