hands-on-modern-rl：动手学强化学习策略梯度reinforce

Cachel wood2026-08-02 10:03

文章目录

策略梯度

单步更新的策略梯度方法，比 REINFORCE 更高效

REINFORCE 的问题：

  • 必须等一个完整回合结束才能更新（Monte Carlo 方法）
  • 如果回合很长，数据利用效率低
  • 回合结束后的高方差 G t G_t Gt用来更新前面所有步骤

Actor-Critic 的改进：

  • 用 TD(0) 估计替代完整回合回报
  • a d v a n t a g e = r + γ ∗ V ( s ′ ) − V ( s ) advantage = r + γ * V(s') - V(s) advantage=r+γ∗V(s′)−V(s)
  • 每一步都可以立即更新，不需要等回合结束
  • 因为用 V ( s ′ ) V(s') V(s′) 自举（bootstrap），方差更低

网络结构：

共享骨干层的 A c t o r − C r i t i c Actor-Critic Actor−Critic 网络

  • 共享层：提取状态特征（复用参数，减少计算量）
  • Actor 头：输出动作概率（策略）
  • Critic 头：输出状态价值（价值函数）
python 复制代码 
"""
第5章：Actor-Critic 算法 ------ CartPole-v1
单步更新的策略梯度方法，比 REINFORCE 更高效

REINFORCE 的问题：
    - 必须等一个完整回合结束才能更新（Monte Carlo 方法）
    - 如果回合很长，数据利用效率低
    - 回合结束后的高方差 G_t 用来更新前面所有步骤

Actor-Critic 的改进：
    - 用 TD(0) 估计替代完整回合回报
    - advantage = r + γ * V(s') - V(s)
    - 每一步都可以立即更新，不需要等回合结束
    - 因为用 V(s') 自举（bootstrap），方差更低

网络结构：
    共享骨干层的 Actor-Critic 网络
    - 共享层：提取状态特征（复用参数，减少计算量）
    - Actor 头：输出动作概率（策略）
    - Critic 头：输出状态价值（价值函数）

运行方式：
    python actor_critic_cartpole.py
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import gymnasium as gym
import matplotlib.pyplot as plt

# 创建输出目录
os.makedirs("output", exist_ok=True)

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False


# ==========================================
# 第一部分：Actor-Critic 网络结构
# ==========================================
class ActorCritic(nn.Module):
    """
    Actor-Critic 网络：共享骨干，两个输出头

    架构图：
        输入 state (维度=4)
            │
        ┌───────┐
        │ Linear│ 4 → 128
        │  ReLU │
        └───────┘
            │
        ┌───────┐
        │  Actor │ 128 → 2 → Softmax  （策略：选哪个动作）
        └───────┘
            │
        ┌───────┐
        │ Critic │ 128 → 1            （价值：当前状态值多少）
        └───────┘

    共享骨干的好处：
        - 状态特征只需要计算一次
        - Actor 和 Critic 可以共享底层表示
        - 参数更少，训练更快
    """

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

        # 共享骨干层：提取状态的特征表示
        self.shared_backbone = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
        )

        # Actor 头：输出动作概率分布
        self.actor_head = nn.Sequential(
            nn.Linear(hidden_dim, action_dim),
        )

        # Critic 头：输出状态价值（标量）
        self.critic_head = nn.Sequential(
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, x):
        """
        前向传播，同时输出动作概率和状态价值

        参数：
            x: 状态张量 [batch_size, state_dim]
        返回：
            probs: 动作概率 [batch_size, action_dim]
            value: 状态价值 [batch_size]
        """
        # 共享层提取特征
        features = self.shared_backbone(x)

        # Actor 输出动作概率
        action_logits = self.actor_head(features)
        probs = torch.softmax(action_logits, dim=-1)

        # Critic 输出状态价值
        value = self.critic_head(features).squeeze(-1)

        return probs, value


# ==========================================
# 第二部分：计算 TD 误差和优势
# ==========================================
def compute_advantage(reward, value, next_value, gamma=0.99, done=False):
    """
    计算 TD(0) 优势函数

    TD 优势 = r + γ * V(s') - V(s)

    直觉理解：
        - V(s) 是 Critic 对当前状态的"预测分数"
        - r + γ * V(s') 是"实际获得的奖励 + 对未来的新预测"
        - 两者之差就是"预测误差"：比预期好（正）还是差（负）

    与 REINFORCE 的区别：
        REINFORCE: advantage = G_t（完整回合的累计回报）
        Actor-Critic: advantage = r + γ * V(s') - V(s)（单步 TD 误差）

    参数：
        reward: 即时奖励 r_t
        value: 当前状态价值 V(s_t)
        next_value: 下一状态价值 V(s_{t+1})
        gamma: 折扣因子
        done: 回合是否结束
    返回：
        advantage: TD 优势值
    """
    if done:
        # 回合结束时，没有下一个状态，目标 = r_t
        target = reward
    else:
        # TD 目标：r_t + γ * V(s_{t+1})
        target = reward + gamma * next_value

    advantage = target - value
    return advantage, target


# ==========================================
# 第三部分：主训练循环
# ==========================================
def train():
    """
    Actor-Critic 完整训练流程

    核心区别（与 REINFORCE 对比）：
        REINFORCE：收集完整回合 → 计算所有 G_t → 一次反向传播
        Actor-Critic：每一步都计算 TD 误差 → 立即更新网络

    这意味着 Actor-Critic 可以在线学习（online learning），
    不需要等待回合结束，数据利用效率更高。
    """
    # ---------- 超参数 ----------
    num_episodes = 500
    gamma = 0.99
    learning_rate = 1e-3
    hidden_dim = 128

    # ---------- 初始化 ----------
    env = gym.make("CartPole-v1")
    model = ActorCritic(
        state_dim=env.observation_space.shape[0],
        action_dim=env.action_space.n,
        hidden_dim=hidden_dim,
    )
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # 记录训练数据
    episode_rewards = []
    episode_actor_losses = []
    episode_critic_losses = []

    print("=" * 60)
    print("  Actor-Critic ------ CartPole-v1 训练")
    print("=" * 60)
    print(f"  超参数:")
    print(f"    回合数: {num_episodes}")
    print(f"    折扣因子 γ: {gamma}")
    print(f"    学习率: {learning_rate}")
    print(f"    隐藏层维度: {hidden_dim}")
    print("=" * 60)

    for episode in range(num_episodes):
        state, _ = env.reset()
        episode_reward = 0
        total_actor_loss = 0
        total_critic_loss = 0
        steps = 0

        done = False
        truncated = False

        while not (done or truncated):
            # ========== 第一步：观察当前状态 ==========
            state_tensor = torch.FloatTensor(state).unsqueeze(0)

            # 前向传播：同时获取动作概率和状态价值
            probs, value = model(state_tensor)
            probs = probs.squeeze(0)     # [action_dim]
            value = value.squeeze()       # 标量

            # ========== 第二步：选择动作（按概率采样） ==========
            dist = torch.distributions.Categorical(probs)
            action = dist.sample()

            # 保存 log π(a|s)，用于后续计算策略梯度
            log_prob = dist.log_prob(action)

            # ========== 第三步：执行动作，观察转移 ==========
            next_state, reward, done, truncated, _ = env.step(action.item())
            episode_reward += reward
            steps += 1

            # ========== 第四步：计算下一状态的价值 ==========
            next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0)
            with torch.no_grad():
                _, next_value = model(next_state_tensor)
                next_value = next_value.squeeze()

            # ========== 第五步：计算 TD 优势和损失 ==========
            # TD 优势：A(s,a) = r + γ * V(s') - V(s)
            is_done = done or truncated
            advantage, target = compute_advantage(
                reward, value, next_value, gamma, done=is_done
            )

            # Actor 损失：-log π(a|s) * A(s,a)
            # 与 REINFORCE 形式相同，但 advantage 是单步 TD 估计
            actor_loss = -log_prob * advantage

            # Critic 损失：让 V(s) 逼近 TD 目标 r + γ * V(s')
            critic_loss = nn.MSELoss()(value, target.detach())

            # 合并损失（可以加权，这里等权）
            total_loss = actor_loss + critic_loss

            # ========== 第六步：立即更新网络 ==========
            # 注意：REINFORCE 是回合结束后才更新，这里是每一步都更新！
            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()

            total_actor_loss += actor_loss.item()
            total_critic_losses_save = critic_loss.item()
            total_critic_loss += total_critic_losses_save

            # 移到下一个状态
            state = next_state

        # 记录本回合数据
        episode_rewards.append(episode_reward)
        episode_actor_losses.append(total_actor_loss / max(steps, 1))
        episode_critic_losses.append(total_critic_loss / max(steps, 1))

        # 每 50 回合打印进度
        if (episode + 1) % 50 == 0:
            recent_avg = np.mean(episode_rewards[-50:])
            print(
                f"  回合 {episode + 1:4d}/{num_episodes} | "
                f"本轮奖励: {episode_reward:6.1f} | "
                f"近50均值: {recent_avg:6.1f} | "
                f"步数: {steps:3d}"
            )

    env.close()

    # ---------- 训练结果汇总 ----------
    print("=" * 60)
    print("  训练完成！")
    print(f"  最后 50 回合平均奖励: {np.mean(episode_rewards[-50:]):.1f}")
    print(f"  最佳回合奖励: {np.max(episode_rewards):.1f}")
    print("=" * 60)

    # ---------- 与 REINFORCE 收敛速度对比 ----------
    compare_with_reinforce(episode_rewards)

    # ---------- 绘制训练曲线 ----------
    plot_training_curve(episode_rewards)


# ==========================================
# 第四部分：与 REINFORCE 收敛速度对比
# ==========================================
def compare_with_reinforce(actor_critic_rewards):
    """
    对比 Actor-Critic 与 REINFORCE 的收敛速度

    收敛速度衡量标准：第一次达到目标奖励（如 195）需要多少回合
    CartPole-v1 的"解决"标准：连续 100 回合平均奖励 >= 195
    """
    target_reward = 195

    # 计算 Actor-Critic 的收敛速度
    ac_solve_episode = None
    for i in range(len(actor_critic_rewards) - 99):
        window_avg = np.mean(actor_critic_rewards[i:i + 100])
        if window_avg >= target_reward:
            ac_solve_episode = i + 100
            break

    print("\n" + "-" * 60)
    print("  收敛速度对比")
    print("-" * 60)
    print(f"  CartPole-v1 解决标准: 连续100回合平均奖励 >= {target_reward}")

    if ac_solve_episode:
        print(f"  Actor-Critic 在第 {ac_solve_episode} 回合解决环境")
    else:
        print(f"  Actor-Critic 在 {len(actor_critic_rewards)} 回合内未达到解决标准")

    # 关于 REINFORCE 的说明
    print(f"\n  【参考】一般经验值:")
    print(f"    REINFORCE 通常需要 300-500+ 回合才能解决 CartPole")
    print(f"    Actor-Critic 通常在 200-350 回合内解决")
    print(f"    原因：Actor-Critic 单步更新，数据利用效率更高")
    print(f"           TD(0) 的方差比 Monte Carlo 回报更低")
    print("-" * 60)


# ==========================================
# 第五部分：绘制训练曲线
# ==========================================
def plot_training_curve(episode_rewards):
    """
    绘制 Actor-Critic 的训练奖励曲线

    包含原始奖励和滑动平均线
    """
    fig, ax = plt.subplots(figsize=(10, 5))

    # 原始奖励曲线
    ax.plot(episode_rewards, alpha=0.3, color='steelblue', label='回合奖励（原始）')

    # 滑动平均曲线
    window = 50
    moving_avg = [np.mean(episode_rewards[max(0, i - window + 1):i + 1])
                  for i in range(len(episode_rewards))]
    ax.plot(moving_avg, color='crimson', linewidth=2.0,
            label=f'滑动平均（窗口={window}）')

    # 标注解决标准线
    ax.axhline(y=195, color='green', linestyle='--', alpha=0.7,
               label='解决标准（奖励=195）')

    ax.set_xlabel('训练回合', fontsize=12)
    ax.set_ylabel('回合奖励', fontsize=12)
    ax.set_title('Actor-Critic ------ CartPole-v1 训练曲线', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('output/actor_critic_cartpole_rewards.png', dpi=150, bbox_inches='tight')
    print("  训练曲线已保存为 output/actor_critic_cartpole_rewards.png")
    plt.show()


# ==========================================
# 程序入口
# ==========================================
if __name__ == "__main__":
    train()

REINFORCE

从零实现最经典的策略梯度方法，理解"好动作多做，坏动作少做"

算法核心思想：

  • 策略梯度的直观理解 ------ 如果一个回合得分很高，那么这个回合里的每个动作都应该被"鼓励"（增大概率）；

  • 反之则应该被"抑制"（降低概率）。

  • REINFORCE 公式：

    ∇ J ( θ ) ≈ Σ t ∇ log ⁡ π ( a t ∣ s t ) ∗ G t ∇J(θ) ≈ Σ_t ∇\\log π(a_t\|s_t) * G_t ∇J(θ)≈Σt∇logπ(at∣st)∗Gt

    其中 G t G_t Gt 是从时间步 t t t 开始的折扣累计回报

python 复制代码 
"""
第5章：REINFORCE 策略梯度算法 ------ CartPole-v1
从零实现最经典的策略梯度方法，理解"好动作多做，坏动作少做"

算法核心思想：
    策略梯度的直观理解 ------ 如果一个回合得分很高，
    那么这个回合里的每个动作都应该被"鼓励"（增大概率）；
    反之则应该被"抑制"（降低概率）。

REINFORCE 公式：
    ∇J(θ) ≈ Σ_t [∇log π(a_t|s_t)] * G_t
    其中 G_t 是从时间步 t 开始的折扣累计回报

运行方式：
    python reinforce_cartpole.py
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import gymnasium as gym
import matplotlib.pyplot as plt
from collections import deque

# 创建输出目录
os.makedirs("output", exist_ok=True)

# 设置中文字体，确保图表标题和标签正常显示
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False


# ==========================================
# 第一部分：策略网络（Policy Network）
# ==========================================
class PolicyNetwork(nn.Module):
    """
    策略网络：将状态映射为动作概率分布

    结构：4 (状态维度) → 128 → 128 → 2 (动作维度)
    输出经过 Softmax 归一化，得到合法的概率分布

    CartPole 的状态空间：[小车位置, 小车速度, 杆子角度, 杆子角速度]
    CartPole 的动作空间：[向左推, 向右推]
    """

    def __init__(self, state_dim=4, action_dim=2, hidden_dim=128):
        super(PolicyNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),   # 输入层 → 第一个隐藏层
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),   # 第一个隐藏层 → 第二个隐藏层
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),   # 第二个隐藏层 → 输出层（logits）
        )

    def forward(self, x):
        """
        前向传播：状态 → 动作概率

        参数：
            x: 状态张量，形状 [batch_size, state_dim]
        返回：
            probs: 动作概率，形状 [batch_size, action_dim]，已经过 Softmax
        """
        logits = self.network(x)
        probs = torch.softmax(logits, dim=-1)
        return probs


# ==========================================
# 第二部分：计算折扣累计回报（Returns）
# ==========================================
def compute_returns(rewards, gamma=0.99):
    """
    从后向前计算每一步的折扣累计回报 G_t

    公式：G_t = r_t + γ * r_{t+1} + γ² * r_{t+2} + ...

    示例（gamma=0.99）：
        rewards = [1, 1, 1, 1, 1]
        G_0 = 1 + 0.99*1 + 0.99²*1 + ... ≈ 4.90
        G_4 = 1

    参数：
        rewards: 每一步的即时奖励列表
        gamma: 折扣因子，越接近1越重视未来奖励
    返回：
        returns: 每一步的折扣累计回报列表
    """
    returns = []
    G = 0  # 累计回报

    # 从后向前遍历：利用 G_t = r_t + gamma * G_{t+1} 的递推关系
    for reward in reversed(rewards):
        G = reward + gamma * G
        returns.insert(0, G)  # 在列表头部插入，保持时间顺序

    return returns


# ==========================================
# 第三部分：收集完整回合轨迹
# ==========================================
def collect_episode(policy, env):
    """
    让策略网络在环境中完成一个完整回合，收集轨迹数据

    REINFORCE 是 on-policy 算法，必须用当前策略收集数据，
    用完即丢弃，下一轮需要重新收集。

    参数：
        policy: 策略网络
        env: Gymnasium 环境
    返回：
        states: 状态列表
        actions: 动作列表
        rewards: 奖励列表
        episode_reward: 回合总奖励
    """
    state, _ = env.reset()
    states, actions, rewards = [], [], []

    done = False
    truncated = False

    while not (done or truncated):
        # 将状态转为张量
        state_tensor = torch.FloatTensor(state).unsqueeze(0)  # 添加 batch 维度

        # 获取动作概率分布
        with torch.no_grad():
            probs = policy(state_tensor)

        # 按概率分布采样动作（探索的关键！不是取 argmax）
        dist = torch.distributions.Categorical(probs)
        action = dist.sample().item()

        # 执行动作，观察结果
        next_state, reward, done, truncated, _ = env.step(action)

        # 存储转移数据
        states.append(state)
        actions.append(action)
        rewards.append(reward)

        state = next_state

    episode_reward = sum(rewards)
    return states, actions, rewards, episode_reward


# ==========================================
# 第四部分：训练一个回合（REINFORCE 核心更新）
# ==========================================
def train_one_episode(policy, optimizer, states, actions, returns):
    """
    REINFORCE 的核心：用策略梯度公式更新网络参数

    损失函数 = - Σ_t [log π(a_t|s_t) * G_t]

    这个损失函数的梯度恰好等于策略梯度：
    ∇loss = - Σ_t [∇log π(a_t|s_t) * G_t] = -∇J(θ)

    所以 minimize loss = maximize J(θ)（期望回报）

    参数：
        policy: 策略网络
        optimizer: 优化器
        states: 状态列表
        actions: 动作列表
        returns: 折扣累计回报列表
    返回：
        loss_value: 本轮损失值
    """
    # 将数据转为张量
    states_tensor = torch.FloatTensor(np.array(states))
    actions_tensor = torch.LongTensor(actions)
    returns_tensor = torch.FloatTensor(returns)

    # 前向传播：获取每个状态下的动作概率
    probs = policy(states_tensor)

    # 计算所采取动作的对数概率 log π(a_t|s_t)
    # gather(1, actions) 选取每个状态对应动作的概率
    action_probs = probs.gather(1, actions_tensor.unsqueeze(1)).squeeze(1)
    log_probs = torch.log(action_probs + 1e-8)  # 加小常数防止 log(0)

    # 策略梯度损失：-log π(a_t|s_t) * G_t
    # 直觉理解：
    #   如果 G_t > 0（好结果），-log_prob * G_t < 0，梯度下降会增大 log_prob → 增大概率
    #   如果 G_t < 0（坏结果），-log_prob * G_t > 0，梯度下降会减小 log_prob → 降低概率
    loss = -(log_probs * returns_tensor).mean()

    # 反向传播 + 参数更新
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss.item()


# ==========================================
# 第五部分：主训练循环
# ==========================================
def train():
    """
    REINFORCE 完整训练流程

    超参数说明：
        - num_episodes = 500：训练 500 个回合
        - gamma = 0.99：折扣因子，重视长期回报
        - learning_rate = 1e-3：学习率
        - hidden_dim = 128：隐藏层宽度
    """
    # ---------- 超参数 ----------
    num_episodes = 500
    gamma = 0.99
    learning_rate = 1e-3
    hidden_dim = 128

    # ---------- 初始化 ----------
    env = gym.make("CartPole-v1")
    policy = PolicyNetwork(
        state_dim=env.observation_space.shape[0],
        action_dim=env.action_space.n,
        hidden_dim=hidden_dim,
    )
    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)

    # 记录训练过程
    episode_rewards = []  # 每个回合的总奖励
    episode_losses = []   # 每个回合的损失

    print("=" * 60)
    print("  REINFORCE 策略梯度 ------ CartPole-v1 训练")
    print("=" * 60)
    print(f"  超参数:")
    print(f"    回合数: {num_episodes}")
    print(f"    折扣因子 γ: {gamma}")
    print(f"    学习率: {learning_rate}")
    print(f"    隐藏层维度: {hidden_dim}")
    print("=" * 60)

    # ---------- 训练循环 ----------
    for episode in range(num_episodes):
        # 第一步：用当前策略收集一个完整回合的轨迹
        states, actions, rewards, episode_reward = collect_episode(policy, env)

        # 第二步：计算折扣累计回报
        returns = compute_returns(rewards, gamma=gamma)

        # 第三步：执行策略梯度更新
        loss_value = train_one_episode(policy, optimizer, states, actions, returns)

        # 记录数据
        episode_rewards.append(episode_reward)
        episode_losses.append(loss_value)

        # 每 50 个回合打印一次进度
        if (episode + 1) % 50 == 0:
            recent_rewards = episode_rewards[-50:]
            avg_reward = np.mean(recent_rewards)
            print(
                f"  回合 {episode + 1:4d}/{num_episodes} | "
                f"本轮奖励: {episode_reward:6.1f} | "
                f"近 50 回合均值: {avg_reward:6.1f} | "
                f"损失: {loss_value:.4f}"
            )

    env.close()

    # ---------- 训练结果汇总 ----------
    print("=" * 60)
    print("  训练完成！")
    print(f"  最后 50 回合平均奖励: {np.mean(episode_rewards[-50:]):.1f}")
    print(f"  最佳回合奖励: {np.max(episode_rewards):.1f}")
    print("=" * 60)

    # ---------- 绘制训练曲线 ----------
    plot_training_curve(episode_rewards)


# ==========================================
# 第六部分：绘制训练曲线
# ==========================================
def plot_training_curve(episode_rewards):
    """
    绘制奖励曲线和滑动平均线

    滑动平均（window=50）可以更清晰地展示学习趋势，
    过滤掉单回合的随机波动。
    """
    fig, ax = plt.subplots(figsize=(10, 5))

    # 原始奖励曲线（浅色，展示波动）
    ax.plot(episode_rewards, alpha=0.3, color='steelblue', label='回合奖励（原始）')

    # 滑动平均曲线（深色，展示趋势）
    window = 50
    if len(episode_rewards) >= window:
        moving_avg = []
        for i in range(len(episode_rewards)):
            start = max(0, i - window + 1)
            moving_avg.append(np.mean(episode_rewards[start:i + 1]))
        ax.plot(moving_avg, color='crimson', linewidth=2.0,
                label=f'滑动平均（窗口={window}）')

    ax.set_xlabel('训练回合', fontsize=12)
    ax.set_ylabel('回合奖励', fontsize=12)
    ax.set_title('REINFORCE 策略梯度 ------ CartPole-v1 训练曲线', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('output/reinforce_cartpole_rewards.png', dpi=150, bbox_inches='tight')
    print("  训练曲线已保存为 output/reinforce_cartpole_rewards.png")
    plt.show()


# ==========================================
# 程序入口
# ==========================================
if __name__ == "__main__":
    train()

带baseline的reinforce

对比原始 REINFORCE 与加入基线（Value Network）的版本

核心问题：原始 REINFORCE 的梯度方差很大，训练不稳定

解决方案：用优势函数代替原始回报

优势 = G t − V ( s t ) G_t - V(s_t) Gt−V(st)

其中 V ( s t ) V(s_t) V(st) 是一个价值网络对状态的估计值

为什么基线能降低方差？

  • G t G_t Gt 的绝对值可能很大（比如 200），但不同时间步的差异较小

  • 减去 V ( s ) V(s) V(s) 后，优势值围绕 0 波动，幅度小得多

  • 数学上 E G t − b = E G t EG_t - b = EG_t EGt−b=EGt（只要 b b b 不依赖动作），期望不变，方差降低

python 复制代码 
"""
第5章：REINFORCE with Baseline ------ 方差缩减对比实验
对比原始 REINFORCE 与加入基线（Value Network）的版本

核心问题：原始 REINFORCE 的梯度方差很大，训练不稳定
解决方案：用优势函数代替原始回报
    优势 = G_t - V(s_t)
    其中 V(s_t) 是一个价值网络对状态的估计值

为什么基线能降低方差？
    - G_t 的绝对值可能很大（比如 200），但不同时间步的差异较小
    - 减去 V(s) 后，优势值围绕 0 波动，幅度小得多
    - 数学上 E[G_t - b] = E[G_t]（只要 b 不依赖动作），期望不变，方差降低

运行方式：
    python reinforce_with_baseline.py
"""

import os
import random
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import gymnasium as gym
import matplotlib.pyplot as plt

# 创建输出目录
os.makedirs("output", exist_ok=True)
SEED = 0

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False


# ==========================================
# 第一部分：网络结构定义
# ==========================================
class PolicyNetwork(nn.Module):
    """
    策略网络（Actor）：状态 → 动作概率

    结构：4 → 128 → 128 → 2（Softmax 输出）
    """

    def __init__(self, state_dim=4, action_dim=2, hidden_dim=128):
        super(PolicyNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
        )

    def forward(self, x):
        logits = self.network(x)
        probs = torch.softmax(logits, dim=-1)
        return probs


class ValueNetwork(nn.Module):
    """
    价值网络（Baseline/Critic）：状态 → 价值估计

    结构：4 → 128 → 128 → 1（标量输出）
    用于估计 V(s)，即从状态 s 出发的期望累计回报

    这个网络就是"基线"：通过减去 V(s)，我们得到优势函数 A(s)
    """

    def __init__(self, state_dim=4, hidden_dim=128):
        super(ValueNetwork, self).__init__()
        self.network = 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, x):
        return self.network(x).squeeze(-1)  # 去掉最后一维，变为 [batch_size]


# ==========================================
# 第二部分：计算折扣累计回报
# ==========================================
def compute_returns(rewards, gamma=0.99):
    """
    计算折扣累计回报 G_t = r_t + γ * r_{t+1} + γ² * r_{t+2} + ...

    参数：
        rewards: 即时奖励列表
        gamma: 折扣因子
    返回：
        returns: 折扣累计回报列表
    """
    returns = []
    G = 0
    for reward in reversed(rewards):
        G = reward + gamma * G
        returns.insert(0, G)
    return returns


# ==========================================
# 第三部分：收集回合轨迹
# ==========================================
def collect_episode(policy, env):
    """
    用当前策略收集一个完整回合的数据

    参数：
        policy: 策略网络
        env: 环境
    返回：
        states, actions, rewards, episode_reward
    """
    state, _ = env.reset()
    states, actions, rewards = [], [], []
    done, truncated = False, False

    while not (done or truncated):
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            probs = policy(state_tensor)
        dist = torch.distributions.Categorical(probs)
        action = dist.sample().item()

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

        states.append(state)
        actions.append(action)
        rewards.append(reward)
        state = next_state

    episode_reward = sum(rewards)
    return states, actions, rewards, episode_reward


# ==========================================
# 第四部分：原始 REINFORCE 训练
# ==========================================
def train_vanilla_reinforce(num_episodes=500, gamma=0.99, lr=1e-3):
    """
    原始 REINFORCE（无基线）

    损失 = -Σ log π(a_t|s_t) * G_t
    直接用折扣累计回报 G_t 作为权重
    """
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)

    env = gym.make("CartPole-v1")
    env.reset(seed=SEED)
    policy = PolicyNetwork(
        state_dim=env.observation_space.shape[0],
        action_dim=env.action_space.n,
    )
    optimizer = optim.Adam(policy.parameters(), lr=lr)

    episode_rewards = []
    gradient_estimates = []  # 记录梯度估计值，用于衡量方差

    for episode in range(num_episodes):
        # 收集轨迹
        states, actions, rewards, episode_reward = collect_episode(policy, env)

        # 计算回报
        returns = compute_returns(rewards, gamma)

        # 转为张量
        states_t = torch.FloatTensor(np.array(states))
        actions_t = torch.LongTensor(actions)
        returns_t = torch.FloatTensor(returns)

        # 前向传播
        probs = policy(states_t)
        action_probs = probs.gather(1, actions_t.unsqueeze(1)).squeeze(1)
        log_probs = torch.log(action_probs + 1e-8)

        # 策略梯度损失
        loss = -(log_probs * returns_t).mean()

        # 记录梯度估计值（用于后续计算方差）
        with torch.no_grad():
            grad_estimate = (log_probs * returns_t).mean().item()
            gradient_estimates.append(grad_estimate)

        # 更新
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        episode_rewards.append(episode_reward)

        if (episode + 1) % 100 == 0:
            avg = np.mean(episode_rewards[-50:])
            print(f"  [Vanilla] 回合 {episode+1:4d} | 近50均值: {avg:6.1f}")

    env.close()
    return episode_rewards, gradient_estimates


# ==========================================
# 第五部分：REINFORCE with Baseline 训练
# ==========================================
def train_reinforce_with_baseline(num_episodes=500, gamma=0.99, lr=1e-3):
    """
    REINFORCE + 价值基线

    优势函数：A(s,a) = G_t - V(s_t)
    策略损失：-Σ log π(a_t|s_t) * A(s_t, a_t)
    价值损失：MSE(V(s_t), G_t)

    两个网络同时训练：
        - 策略网络学习"什么动作更好"（相对于基线）
        - 价值网络学习"当前状态平均能拿多少分"（基线）
    """
    baseline_seed = SEED + 100
    random.seed(baseline_seed)
    np.random.seed(baseline_seed)
    torch.manual_seed(baseline_seed)

    env = gym.make("CartPole-v1")
    env.reset(seed=baseline_seed)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    # 初始化策略网络和价值网络
    policy = PolicyNetwork(state_dim=state_dim, action_dim=action_dim)
    value_net = ValueNetwork(state_dim=state_dim)

    # 两个网络各用独立的优化器
    policy_optimizer = optim.Adam(policy.parameters(), lr=lr)
    value_optimizer = optim.Adam(value_net.parameters(), lr=lr)

    episode_rewards = []
    gradient_estimates = []  # 记录梯度估计值

    for episode in range(num_episodes):
        # 收集轨迹
        states, actions, rewards, episode_reward = collect_episode(policy, env)

        # 计算回报
        returns = compute_returns(rewards, gamma)

        # 转为张量
        states_t = torch.FloatTensor(np.array(states))
        actions_t = torch.LongTensor(actions)
        returns_t = torch.FloatTensor(returns)

        # ========== 更新价值网络（Critic） ==========
        # 价值网络的目标：准确预测 V(s) ≈ G_t
        values = value_net(states_t)
        value_loss = nn.MSELoss()(values, returns_t)

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

        # ========== 计算优势函数 ==========
        # 优势 = 实际回报 - 基线预测
        # A > 0 表示"比预期好" → 增大对应动作概率
        # A < 0 表示"比预期差" → 减小对应动作概率
        with torch.no_grad():
            values_pred = value_net(states_t)
        advantages = returns_t - values_pred

        # ========== 更新策略网络（Actor） ==========
        probs = policy(states_t)
        action_probs = probs.gather(1, actions_t.unsqueeze(1)).squeeze(1)
        log_probs = torch.log(action_probs + 1e-8)

        # 策略梯度损失：用优势函数替代原始回报
        policy_loss = -(log_probs * advantages).mean()

        # 记录梯度估计值（用优势替代回报）
        with torch.no_grad():
            grad_estimate = (log_probs * advantages).mean().item()
            gradient_estimates.append(grad_estimate)

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

        episode_rewards.append(episode_reward)

        if (episode + 1) % 100 == 0:
            avg = np.mean(episode_rewards[-50:])
            print(f"  [Value Baseline] 回合 {episode+1:4d} | 近50均值: {avg:6.1f}")

    env.close()
    return episode_rewards, gradient_estimates


# ==========================================
# 第六部分：对比实验主函数
# ==========================================
def run_comparison():
    """
    运行对比实验：Vanilla REINFORCE vs REINFORCE + Value Baseline

    对比两个维度：
        1. 学习速度和最终性能（奖励曲线）
        2. 梯度估计的方差（方差越低，训练越稳定）
    """
    num_episodes = 500
    gamma = 0.99
    lr = 1e-3

    print("=" * 60)
    print("  REINFORCE 方差缩减对比实验")
    print("=" * 60)
    print(f"  训练回合数: {num_episodes}")
    print(f"  折扣因子 γ: {gamma}")
    print(f"  学习率: {lr}")
    print("=" * 60)

    # ---------- 实验1：原始 REINFORCE ----------
    print("\n[实验1] 训练 Vanilla REINFORCE（无基线）...")
    vanilla_rewards, vanilla_grads = train_vanilla_reinforce(
        num_episodes=num_episodes, gamma=gamma, lr=lr
    )

    # ---------- 实验2：REINFORCE + Value Baseline ----------
    print("\n[实验2] 训练 REINFORCE + Value Baseline（价值基线）...")
    baseline_rewards, baseline_grads = train_reinforce_with_baseline(
        num_episodes=num_episodes, gamma=gamma, lr=lr
    )

    # ---------- 方差统计对比 ----------
    print("\n" + "=" * 60)
    print("  方差对比统计")
    print("=" * 60)

    vanilla_grad_var = np.var(vanilla_grads)
    baseline_grad_var = np.var(baseline_grads)

    print(f"  Vanilla REINFORCE 梯度估计方差: {vanilla_grad_var:.6f}")
    print(f"  REINFORCE+Value Baseline 梯度估计方差: {baseline_grad_var:.6f}")

    if vanilla_grad_var > 0:
        ratio = vanilla_grad_var / max(baseline_grad_var, 1e-10)
        print(f"  方差比（Vanilla/Value Baseline）: {ratio:.2f}x")
        print(f"  Value Baseline 将方差降低至原来的 {1/ratio*100:.1f}%")

    print(f"\n  Vanilla REINFORCE 最后50回合均值: {np.mean(vanilla_rewards[-50:]):.1f}")
    print(f"  REINFORCE+Value Baseline 最后50回合均值: {np.mean(baseline_rewards[-50:]):.1f}")
    print("=" * 60)

    # ---------- 绘制对比图1：奖励曲线 ----------
    plot_reward_comparison(vanilla_rewards, baseline_rewards, window=50)

    # ---------- 绘制对比图2：方差对比 ----------
    plot_variance_comparison(vanilla_grads, baseline_grads, window=50)


# ==========================================
# 第七部分：绘制奖励对比曲线
# ==========================================
def plot_reward_comparison(vanilla_rewards, baseline_rewards, window=50):
    """
    绘制两组实验的奖励曲线对比

    包含原始曲线和滑动平均曲线，直观展示：
        - Value Baseline 版本是否收敛更快
        - Value Baseline 版本是否更稳定（波动更小）
    """
    fig, ax = plt.subplots(figsize=(10, 5))

    # Vanilla REINFORCE
    ax.plot(vanilla_rewards, alpha=0.2, color='steelblue')
    vanilla_avg = [np.mean(vanilla_rewards[max(0, i-window+1):i+1])
                   for i in range(len(vanilla_rewards))]
    ax.plot(vanilla_avg, color='steelblue', linewidth=2.0,
            label='Vanilla REINFORCE')

    # REINFORCE + Value Baseline
    ax.plot(baseline_rewards, alpha=0.2, color='crimson')
    baseline_avg = [np.mean(baseline_rewards[max(0, i-window+1):i+1])
                    for i in range(len(baseline_rewards))]
    ax.plot(baseline_avg, color='crimson', linewidth=2.0,
            label='REINFORCE + Value Baseline')

    ax.set_xlabel('训练回合', fontsize=12)
    ax.set_ylabel('回合奖励', fontsize=12)
    ax.set_title('REINFORCE 奖励曲线对比（Vanilla vs Value Baseline）', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('output/reinforce_baseline_reward_comparison.png', dpi=150, bbox_inches='tight')
    print("  奖励对比图已保存为 output/reinforce_baseline_reward_comparison.png")
    plt.show()


# ==========================================
# 第八部分：绘制方差对比图
# ==========================================
def plot_variance_comparison(vanilla_grads, baseline_grads, window=50):
    """
    绘制梯度估计方差的滑动窗口对比

    这张图是本实验的核心：展示 Value Baseline 如何降低策略梯度的方差。
    方差越低，训练过程越稳定，收敛越可靠。
    """
    fig, ax = plt.subplots(figsize=(10, 5))

    # 计算滑动窗口方差
    def moving_variance(data, w):
        variances = []
        for i in range(len(data)):
            start = max(0, i - w + 1)
            variances.append(np.var(data[start:i + 1]))
        return variances

    vanilla_var = moving_variance(vanilla_grads, window)
    baseline_var = moving_variance(baseline_grads, window)

    ax.plot(vanilla_var, color='steelblue', linewidth=1.5, alpha=0.8,
            label='Vanilla REINFORCE')
    ax.plot(baseline_var, color='crimson', linewidth=1.5, alpha=0.8,
            label='REINFORCE + Value Baseline')

    ax.set_xlabel('训练回合', fontsize=12)
    ax.set_ylabel(f'梯度估计方差（窗口={window}）', fontsize=12)
    ax.set_title('策略梯度方差对比 ------ Value Baseline 的方差缩减效果', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)

    # 添加注释箭头，标明方差差异
    if len(vanilla_var) > 100:
        mid_point = len(vanilla_var) // 2
        ax.annotate(
            'Value Baseline 降低方差',
            xy=(mid_point, baseline_var[mid_point]),
            xytext=(mid_point + 50, max(vanilla_var) * 0.7),
            fontsize=11,
            arrowprops=dict(arrowstyle='->', color='gray'),
            color='gray',
        )

    plt.tight_layout()
    plt.savefig('output/reinforce_baseline_variance_comparison.png', dpi=150, bbox_inches='tight')
    print("  方差对比图已保存为 output/reinforce_baseline_variance_comparison.png")
    plt.show()


# ==========================================
# 程序入口
# ==========================================
if __name__ == "__main__":
    run_comparison()
相关推荐
蓝斯4971 小时前
一碰即传，重构跨设备文件分享体验
开发语言·python·重构
宁风NF1 小时前
JavaScript：内存、垃圾回收、性能优化
开发语言·前端·javascript·学习·性能优化·es6
寒水馨1 小时前
macOS下载、安装uv-0.12.0（附安装包uv-aarch64-apple-darwin.tar.gz）
python·macos·rust·项目管理·包管理器·astral·pip替代
ShuiShenHuoLe2 小时前
Go html/template 使用入门
开发语言·golang·html
geovindu2 小时前
java: Gale-Shapley Algorithm
java·开发语言·后端·算法
冻柠檬飞冰走茶2 小时前
PTA基础编程题目集 7-34 通讯录的录入与显示（C语言实现）
c语言·开发语言·数据结构·算法
星核0penstarry2 小时前
DeepSeek-V4-Flash 正式公测：大模型行业进入「极速平价普惠时代」
java·开发语言·人工智能
热门推荐
01如何新建文件夹？ 电脑新建文件夹的4种方法02GitHub 镜像站点03国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）042026年7月AI圈大地震：GPT-5.6被政府限制、Claude入驻Slack、Anthropic自研芯片052026 国产 AI 大模型横评：DeepSeek、通义千问、Kimi、文心一言、星火、豆包谁更能打？06AI 编程 IDE 全景解析 2026：Agent 全面接管开发链路07微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新08AI科技热点日报 | 2026年07月01日09Agnes AI 免费 API 接入指南：文本、生图、生视频，一套接口全免费102026上半年大模型全景技术解读：推理融合、Agent 爆发与多模态统一