策略梯度（Policy Gradient）

第一部分：强化学习基础

在讲解策略梯度之前，我们需要先理解强化学习的基本框架，因为策略梯度是强化学习中的一种方法。

1.1 什么是强化学习？

强化学习是一种机器学习范式，目标是让智能体（Agent）通过与环境（Environment）的交互，学习如何采取行动以最大化长期的累积奖励（Cumulative Reward）。与监督学习（给定输入-输出对）不同，强化学习没有明确的标签，智能体需要通过试错（Trial-and-Error）来发现哪些行动是“好的”。

强化学习的经典例子：

• 游戏：智能体在游戏中选择动作（如移动、跳跃），环境返回奖励（如得分）或惩罚（如游戏失败）。
• 机器人控制：机器人学习走路，环境提供反馈（如是否摔倒）。

1.2 强化学习的核心组件

强化学习的数学框架可以用以下几个核心概念来描述：

1. 状态（State, $S$）：

   • 表示智能体当前所处的环境情况。例如，在迷宫游戏中，状态可以是智能体的位置。
   • 状态空间 $\mathcal{S}$ 是所有可能状态的集合。

2. 动作（Action, $A$）：

   • 智能体在某个状态下可以采取的行为。例如，在迷宫中，动作可以是“向左走”或“向右走”。
   • 动作空间 $\mathcal{A}$ 是所有可能动作的集合。

3. 奖励（Reward, $R$）：

   • 环境对智能体动作的反馈，是一个标量值。例如，走出迷宫可能得到 $+100$ 的奖励，撞墙可能得到 $-10$。
   • 奖励函数 $R(s,a,s^{\prime })$ 定义了在状态 $s$，采取动作 $a$，转移到状态 $s^{\prime }$ 时获得的奖励。

4. 策略 Wǔ（Policy, $\pi$）：

   • 策略是智能体决定动作的规则，定义了在状态 $s$ 下选择动作 $a$ 的概率。
   • 形式化表示为 $\pi (a|s)$，表示在状态 $s$ 下选择动作 $a$ 的概率。
   • 策略可以是确定性的（总是选择某个动作）或随机性的（按概率分布选择动作）。

5. 环境模型（Environment Model）：

   • 环境决定了状态转移和奖励的规则。通常用状态转移概率 $P(s^{\prime }|s,a)$ 表示从状态 $s$，采取动作 $a$，转移到状态 $s^{\prime }$ 的概率。

6. 回报（Return, $G$）：

   • 回报是智能体从某一时刻开始，未来获得的累积奖励。通常用折扣因子 $\gamma \in [0,1]$ 来平衡短期和长期奖励：
     $G_t=R_{t+1}+\gamma R_{t+2}+{\gamma }^2{R}_{t+3}+\cdots ={\sum }_{k=0}^{\infty }{\gamma }^k{R}_{t+k+1}$
     其中 $\gamma$ 越接近 1，智能体越重视长期奖励。

7. 价值函数（Value Function）：

   • 价值函数评估状态或状态-动作对的“价值”，即在该状态下遵循策略 $\pi$ 能获得的期望回报。
   • 状态价值函数：
     $V^{\pi }(s)={\mathbb{E}}_{\pi }[G_t|S_t=s]$
   • 动作价值函数：
     $Q^{\pi }(s,a)={\mathbb{E}}_{\pi }[G_t|S_t=s,A_t=a]$

1.3 强化学习的目标

强化学习的目标是找到一个最优策略 ${\pi }^{\ast }$，使得智能体在任意状态下都能最大化期望累积回报：
${\pi }^{\ast }=\arg {\max }_{\pi }{\mathbb{E}}_{\pi }[G_t]$

1.4 强化学习的分类

强化学习的方法可以分为两类：

• 基于价值的方法（Value-Based Methods）：如 Q-Learning，学习动作价值函数 $Q(s,a)$，然后选择价值最大的动作。
• 基于策略的方法（Policy-Based Methods）：直接优化策略 $\pi (a|s)$，使之输出更好的动作分布。策略梯度就属于这一类。
• 演员-评论家方法（Actor-Critic Methods）：结合基于价值和基于策略的方法，既学习策略（演员），又学习价值函数（评论家）。

---

第二部分：策略梯度简介

2.1 为什么需要策略梯度？

在强化学习中，基于价值的方法（如 Q-Learning）通过估计每个状态-动作对的价值来选择动作。然而，这种方法有以下局限性：

• 离散动作空间：当动作空间很大或连续时，计算每个动作的 $Q(s,a)$ 变得困难。
• 策略的表达能力：基于价值的方法通常选择最大价值的动作（确定性策略），难以处理需要随机性的场景（如部分可观测环境）。
• 复杂环境：在高维或连续动作空间中，基于价值的方法可能收敛缓慢或不稳定。

策略梯度方法通过直接优化策略 $\pi (a|s;\theta )$，克服了这些问题。策略通常用神经网络参数化，参数为 $\theta$，通过梯度上升优化策略，使其输出更可能带来高回报的动作。

2.2 策略梯度的核心思想

策略梯度方法的核心是：

• 将策略 $\pi (a|s;\theta )$ 表示为一个参数化的概率分布（通常用神经网络实现）。
• 定义一个目标函数 $J(\theta )$，表示策略的期望累积回报。
• 通过计算目标函数对参数 $\theta$ 的梯度 ${\nabla }_{\theta }J(\theta )$，使用梯度上升更新参数：
  $\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta )$
  其中 $\alpha$ 是学习率。

这个过程类似于深度学习中的梯度下降，但这里是梯度上升，因为我们希望最大化期望回报。

---

第三部分：策略梯度的数学推导

为了深入理解策略梯度，我们需要推导其核心公式：策略梯度定理。这一部分会涉及一些数学，但我会尽量分解步骤，确保清晰。

3.1 目标函数的定义

我们希望优化的目标是期望累积回报：
$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]$
其中：

• $\tau =(s_0,a_0,r_1,s_1,a_1,r_2,\dots )$ 是一个轨迹（Trajectory），表示智能体与环境的交互序列。
• $R(\tau )={\sum }_{t=0}^T{\gamma }^t{r}_{t+1}$ 是轨迹的累积奖励。
• $\tau \sim {\pi }_{\theta }$ 表示轨迹的概率分布由策略 ${\pi }_{\theta }$ 决定。

轨迹的概率为：
$p(\tau |\theta )=p(s_0){\prod }_{t=0}^T\pi (a_t|s_t;\theta )P(s_{t+1}|s_t,a_t)$
其中：

• $p(s_0)$ 是初始状态分布。
• $\pi (a_t|s_t;\theta )$ 是策略输出的动作概率。
• $P(s_{t+1}|s_t,a_t)$ 是环境的状态转移概率。

因此，目标函数可以写为：
$J(\theta )=\int p(\tau |\theta )R(\tau )d\tau$

3.2 计算策略梯度

我们需要计算目标函数 $J(\theta )$ 对策略参数 $\theta$ 的梯度 ${\nabla }_{\theta }J(\theta )$，以通过梯度上升优化策略 ${\pi }_{\theta }$。目标函数定义为在策略 ${\pi }_{\theta }$ 下轨迹 $\tau$ 的期望回报：$$J(\theta )={\mathbb{E}}_{\tau \sim p(\tau |\theta )}[R(\tau )]=\int p(\tau |\theta )R(\tau )d\tau$$，其中，$\tau =(s_0,a_0,s_1,a_1,\dots ,s_T,a_T,s_{T+1})$ 是一条完整的轨迹，$p(\tau |\theta )$ 是轨迹 $\tau$ 的概率密度，$$R(\tau )=\sum _{t=0}^T{\gamma }^{t}r(s_t,a_t,s_{t+1})$$ 是轨迹的累积回报，$\gamma$ 为折扣因子。

梯度计算

对目标函数求梯度：$${\nabla }_{\theta }J(\theta )={\nabla }_{\theta }\int p(\tau |\theta )R(\tau )d\tau$$
由于积分内部的 $p(\tau |\theta )$ 依赖于 $\theta$，我们将梯度算子移入积分（假设积分满足交换性条件）：$${\nabla }_{\theta }J(\theta )=\int {\nabla }_{\theta }[p(\tau |\theta )R(\tau )]d\tau$$
注意到，$R(\tau )$ 是环境的回报函数，与 $\theta$ 无关，因此：$${\nabla }_{\theta }[p(\tau |\theta )R(\tau )]=R(\tau ){\nabla }_{\theta }p(\tau |\theta )$$
代入：$${\nabla }_{\theta }J(\theta )=\int R(\tau ){\nabla }_{\theta }p(\tau |\theta )d\tau$$

使用对数梯度技巧

为了计算 ${\nabla }_{\theta }p(\tau |\theta )$，我们引入对数梯度技巧（Log-Likelihood Trick）。对于概率密度 $p(\tau |\theta )$，有：$${\nabla }_{\theta }\log p(\tau |\theta )=\frac{{\nabla }_{\theta }p(\tau |\theta )}{p(\tau |\theta )}$$
重排得：$${\nabla }_{\theta }p(\tau |\theta )=p(\tau |\theta ){\nabla }_{\theta }\log p(\tau |\theta )$$
将此代入梯度公式：$${\nabla }_{\theta }J(\theta )=\int R(\tau )[p(\tau |\theta ){\nabla }_{\theta }\log p(\tau |\theta )]d\tau$$
注意到积分 $\int p(\tau |\theta )[\cdot ]d\tau$ 是对轨迹 $\tau \sim p(\tau |\theta )$ 的期望，因此：$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p(\tau |\theta )}[R(\tau ){\nabla }_{\theta }\log p(\tau |\theta )]$$
这一形式表明，策略梯度可以通过采样轨迹 $\tau$ 来估计，具体为轨迹回报 $R(\tau )$ 和轨迹概率对数梯度 ${\nabla }_{\theta }\log p(\tau |\theta )$ 的乘积的期望。

计算轨迹概率的对数梯度

为了计算 ${\nabla }_{\theta }\log p(\tau |\theta )$，需要明确轨迹概率 $p(\tau |\theta )$ 的表达式。轨迹 $\tau$ 的概率由以下部分组成：

• 初始状态分布：$p(s_0)$。
• 策略概率：$\pi (a_t|s_t;\theta )$，表示在状态 $s_t$ 下选择动作 $a_t$ 的概率。
• 状态转移概率：$P(s_{t+1}|s_t,a_t)$，表示环境从状态 $s_t$ 采取动作 $a_t$ 后转移到状态 $s_{t+1}$ 的概率。

轨迹的联合概率为：$$p(\tau |\theta )=p(s_0)\prod _{t=0}^T\pi (a_t|s_t;\theta )P(s_{t+1}|s_t,a_t)$$
取对数：$$\log p(\tau |\theta )=\log p(s_0)+\sum _{t=0}^T\left[\log \pi (a_t|s_t;\theta )+\log P(s_{t+1}|s_t,a_t)\right]$$
对 $\theta$ 求梯度：$${\nabla }_{\theta }\log p(\tau |\theta )={\nabla }_{\theta }\left[\log p(s_0)+\sum _{t=0}^T\left(\log \pi (a_t|s_t;\theta )+\log P(s_{t+1}|s_t,a_t)\right)\right]$$
分析每一项：

• $\log p(s_0)$：初始状态分布由环境决定，不依赖 $\theta$，因此 $${\nabla }_{\theta }\log p(s_0)=0$$
• $\log P(s_{t+1}|s_t,a_t)$：状态转移概率是环境动态模型，不依赖 $\theta$，因此 $${\nabla }_{\theta }\log P(s_{t+1}|s_t,a_t)=0$$
• $\log \pi (a_t|s_t;\theta )$：策略概率显式依赖 $\theta$，其梯度非零。

因此：$${\nabla }_{\theta }\log p(\tau |\theta )=\sum _{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )$$

最终策略梯度表达式

将 ${\nabla }_{\theta }\log p(\tau |\theta )$ 代回：$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p(\tau |\theta )}\left[R(\tau )\sum _{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )\right]$$。这表明策略梯度可以通过以下步骤估计：

1. 从策略 ${\pi }_{\theta }$ 采样多条轨迹 ${\tau }^{(i)},i=1,\dots ,N$。
2. 对每条轨迹，计算累积回报 $R({\tau }^{(i)})$。
3. 计算每个时间步的策略对数概率梯度 ${\nabla }_{\theta }\log \pi (a_t|s_t;\theta )$。
4. 对所有轨迹求平均，近似期望：$${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\left[R({\tau }^{(i)})\sum _{t=0}^T{\nabla }_{\theta }\log \pi (a_t^{(i)}|s_t^{(i)};\theta )\right]$$

补充说明

1. 对数梯度技巧的作用：将 ${\nabla }_{\theta }p(\tau |\theta )$ 转换为 $$p(\tau |\theta ){\nabla }_{\theta }\log p(\tau |\theta )$$，将梯度计算转化为期望形式，便于蒙特卡洛采样，避免直接处理复杂的概率密度。

2. 因果性改进：动作 $a_t$ 只影响时间 $t$ 及之后的奖励，因此可将梯度改写为：$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p(\tau |\theta )}\left[\sum _{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )\sum _{k=t}^T{\gamma }^{k-t}r(s_k,a_k,s_{k+1})\right]$$这利用了因果性，降低了梯度方差。

3. 基线函数：为进一步降低方差，可引入基线函数 $b(s_t)$（如值函数估计），修改梯度为：$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim p(\tau |\theta )}\left[\sum _{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )(R_t-b(s_t))\right]$$，其中 $$R_t=\sum _{k=t}^T{\gamma }^{k-t}r(s_k,a_k,s_{k+1})$$ 是从时间 $t$ 开始的回报。

3.3 策略梯度定理

策略梯度定理是策略梯度方法的核心，它提供了一种简洁的方式来计算目标函数 $J(\theta )$ 的梯度，避免直接对整个轨迹的回报求导。为了进一步简化计算，我们希望梯度只依赖于每个时间步的动作及其价值，而不是整个轨迹的累积回报。策略梯度定理指出：
$${\nabla }_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )Q^{\pi }(s_t,a_t)\right]$$
其中：

• $\pi (a_t|s_t;\theta )$ 是策略，表示在状态 $s_t$ 下选择动作 $a_t$ 的概率。
• ${\nabla }_{\theta }\log \pi (a_t|s_t;\theta )$ 是策略对数概率对参数 $\theta$ 的梯度。
• $Q^{\pi }(s_t,a_t)={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[{\sum }_{k=t}^T{\gamma }^{k-t}{r}_k\mid s_t,a_t\right]$ 是动作价值函数，表示在状态 $s_t$ 采取动作 $a_t$，随后遵循策略 $\pi$ 的期望回报。

定理的意义

策略梯度定理将梯度分解为两个核心部分：

• 策略梯度 ${\nabla }_{\theta }\log \pi (a_t|s_t;\theta )$：表示参数 $\theta$ 对动作 $a_t$ 概率的敏感性，反映调整 $\theta$ 如何改变策略在状态 $s_t$ 下选择动作 $a_t$ 的倾向。
• 动作价值 $Q^{\pi }(s_t,a_t)$：衡量在状态 $s_t$ 采取动作 $a_t$ 后，遵循策略 $\pi$ 的期望回报，代表该动作的“价值”。

通过这两部分的乘积，定理确保策略朝着选择高价值动作的方向优化：

• 如果 $Q^{\pi }(s_t,a_t)>0$，策略会增加 $\pi (a_t|s_t;\theta )$ 的概率，倾向于更频繁地选择动作 $a_t$。
• 如果 $Q^{\pi }(s_t,a_t)<0$，策略会减少 $\pi (a_t|s_t;\theta )$ 的概率，降低选择动作 $a_t$ 的可能性。
• 动作价值 $Q^{\pi }(s_t,a_t)$ 的绝对值越大，策略调整的幅度越大；价值较高的动作会更显著地增加其概率。

这个机制使策略能够通过梯度上升逐步学习到更优的动作选择策略。

3.4 蒙特卡洛采样

在实际中，我们无法直接计算期望，因此使用蒙特卡洛方法采样多条轨迹，近似梯度：
${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}{\sum }_{i=1}^N{\sum }_{t=0}^T{\nabla }_{\theta }\log \pi (a_t^{(i)}|s_t^{(i)};\theta )G_t^{(i)}$
其中：

• $N$ 是采样轨迹数。
• $G_t^{(i)}$ 是第 $i$ 条轨迹在时间 $t$ 的实际回报，近似 $Q^{\pi }(s_t,a_t)$。

---

第四部分：策略梯度算法

4.1 REINFORCE算法

REINFORCE 是最基本的策略梯度算法，基于蒙特卡洛采样。其伪代码如下：

1. 初始化策略参数 $\theta$。
2. 重复以下步骤：
   • 采样 $N$ 条轨迹 ${\tau }^{(i)}$，每条轨迹包含状态、动作和奖励序列。
   • 对于每条轨迹，计算回报 $G_t^{(i)}$。
   • 计算梯度：
     ${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}{\sum }_{i=1}^N{\sum }_{t=0}^T{\nabla }_{\theta }\log \pi (a_t^{(i)}|s_t^{(i)};\theta )G_t^{(i)}$
   • 更新参数：
     $\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta )$

优点：

• 简单，易于实现。
• 适用于离散和连续动作空间。

缺点：

• 方差高：蒙特卡洛采样的回报 $G_t$ 波动大，导致梯度估计不稳定。
• 收敛缓慢：需要大量样本。

4.2 减小方差的技巧

为了提高策略梯度的稳定性，常用以下技巧：

1. 基线（Baseline）：

   • 在梯度中引入基线 $b(s_t)$，不改变期望，但降低方差：
     ${\nabla }_{\theta }J(\theta )\approx \frac{1}{N}{\sum }_{i=1}^N{\sum }_{t=0}^T{\nabla }_{\theta }\log \pi (a_t^{(i)}|s_t^{(i)};\theta )(G_t^{(i)}-b(s_t^{(i)}))$
   • 常用的基线是状态价值函数 $V^{\pi }(s_t)$，可以用神经网络估计。

2. 优势函数（Advantage Function）：

   • 用优势函数 $A(s_t,a_t)=Q(s_t,a_t)-V(s_t)$ 替代 $Q(s_t,a_t)$，进一步降低方差：
     ${\nabla }_{\theta }J(\theta )\approx \mathbb{E}\left[{\sum }_{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )A(s_t,a_t)\right]$

3. 折扣因子：

   • 在计算回报时，使用折扣因子 $\gamma$ 限制未来奖励的影响，减少方差。

4.3 演员-评论家算法（Actor-Critic）

为了解决 REINFORCE 的高方差问题，演员-评论家方法结合了基于策略和基于价值的方法：

• 演员（Actor）：学习策略 $\pi (a|s;\theta )$，负责选择动作。
• 评论家（Critic）：学习价值函数 $V(s;w)$ 或 $Q(s,a;w)$，评估动作的好坏。

梯度更新为：
${\nabla }_{\theta }J(\theta )\approx {\sum }_{t=0}^T{\nabla }_{\theta }\log \pi (a_t|s_t;\theta )A(s_t,a_t)$
其中 $A(s_t,a_t)$ 由评论家估计，例如：
$A(s_t,a_t)=r_{t+1}+\gamma V(s_{t+1};w)-V(s_t;w)$

优点：

• 实时更新，适合在线学习。
• 方差较低，收敛更快。

缺点：

• 需要同时优化演员和评论家，增加了计算复杂性。

---

第五部分：策略梯度的实际应用

5.1 常见应用场景

策略梯度方法在以下领域表现优异：

• 机器人控制：如机械臂抓取、无人机导航，动作空间通常是连续的。
• 游戏：如 Atari 游戏、围棋，策略梯度可以处理复杂的策略空间。
• 自然语言处理：如对话系统，策略梯度用于优化生成序列的策略。
• 金融：如自动交易，策略梯度用于优化投资组合。

5.2 现代策略梯度算法

基于策略梯度的现代算法进一步改进了性能：

1. TRPO（Trust Region Policy Optimization）：

   • 通过限制策略更新的步长，保证稳定性。
   • 使用信任区域约束优化目标函数。

2. PPO（Proximal Policy Optimization）：

   • 改进 TRPO，简化计算，使用剪切概率比（Clipped Surrogate Objective）。
   • 易于实现，广泛应用于实际任务。

3. SAC（Soft Actor-Critic）：

   • 引入熵正则化，鼓励探索，适合连续动作空间。

---

第六部分：实现策略梯度的代码示例

以下是一个简单的 REINFORCE 算法的 Python 实现，使用 PyTorch 和 Gym 环境（以 CartPole 为例）：

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim

class PolicyNetwork(nn.Module):
    """
    策略网络：基于神经网络实现，用于输出动作概率分布。
    该网络用于REINFORCE算法，采用简单的前馈神经网络结构，包含一个隐藏层。
    """
    def __init__(self, input_dim, hidden_dim, output_dim):
        """
        初始化策略网络。

        参数：
            input_dim (int): 输入状态的维度，例如CartPole中的状态向量维度为4。
            hidden_dim (int): 隐藏层的神经元数量，用于控制网络的容量。
            output_dim (int): 输出层的维度，对应可能动作的数量（例如CartPole的2个动作）。
        """
        super(PolicyNetwork, self).__init__()
        # 定义第一全连接层，将输入状态映射到隐藏层
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        # 定义第二全连接层，将隐藏层映射到动作空间
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        # Softmax层，将动作 logits 转换为概率分布
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x):
        """
        前向传播，计算给定状态下的动作概率分布。

        参数：
            x (torch.Tensor): 输入状态张量，形状为 (batch_size, input_dim) 或 (input_dim,)。

        返回：
            torch.Tensor: 动作的概率分布，形状为 (batch_size, output_dim) 或 (output_dim,)。
        """
        # 通过第一层并应用 ReLU 激活函数，增加非线性
        x = torch.relu(self.fc1(x))
        # 通过第二层，得到动作的 logits
        x = self.fc2(x)
        # 应用 Softmax，转换为概率分布
        return self.softmax(x)

def compute_returns(rewards, gamma=0.99):
    """
    计算折扣回报：基于奖励序列计算每个时间步的累积折扣回报，并进行归一化处理。

    策略梯度方法的核心思想：
    - 策略梯度方法是一种强化学习算法，直接优化策略函数 π(a|s;θ)，通过梯度上升最大化期望回报。
    - 折扣回报（discounted return）用于衡量一个动作序列的长期价值，考虑未来的奖励并通过折扣因子 gamma 衰减。
    - 归一化回报可以减少训练过程中的方差，使梯度更新更稳定。

    参数：
        rewards (list): 单次回合中每个时间步的奖励列表。
        gamma (float): 折扣因子，控制未来奖励的重要性，范围为 [0, 1)，通常接近1。

    返回：
        torch.Tensor: 归一化的折扣回报张量，形状为 (episode_length,)。
    """
    returns = []
    discounted_return = 0
    # 从后向前遍历奖励，计算折扣回报
    for reward in reversed(rewards):
        # 当前时间步的折扣回报 = 当前奖励 + gamma * 后续折扣回报
        discounted_return = reward + gamma * discounted_return
        # 将当前时间步的回报插入到列表头部，保持时间顺序
        returns.insert(0, discounted_return)
    
    # 将回报转换为张量
    returns = torch.tensor(returns, dtype=torch.float32)
    # 归一化处理：减去均值，除以标准差（加小值防止除零）
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)
    return returns

def train_reinforce():
    """
    使用REINFORCE算法训练策略网络，应用于CartPole-v1环境。

    REINFORCE算法（Monte Carlo Policy Gradient）的核心思想：
    1. 策略网络 π(a|s;θ) 输出动作概率分布，参数为 θ。
    2. 通过与环境交互，采样一条完整轨迹（状态、动作、奖励序列）。
    3. 计算每个时间步的折扣回报 G_t，作为动作价值的估计。
    4. 计算策略梯度：∇_θ J(θ) ≈ Σ_t [ G_t * ∇_θ log π(a_t|s_t;θ) ]。
    5. 使用梯度上升更新策略网络参数：θ ← θ + α * ∇_θ J(θ)。

    优点：
    - 直接优化策略，适合离散动作空间。
    - 简单直观，易于实现。
    缺点：
    - 基于蒙特卡洛采样的高方差，可能导致训练不稳定。
    - 需要完整轨迹，计算效率较低。

    CartPole-v1 环境：
    - 任务：控制小车移动，使杆保持平衡。
    - 状态：4维向量（小车位置、速度、杆角度、角速度）。
    - 动作：2种（向左或向右推小车）。
    - 奖励：每一步+1，直到回合结束（杆倒下或达到最大时间步）。
    """
    # 初始化环境
    env = gym.make('CartPole-v1')
    input_dim = env.observation_space.shape[0]  # 状态维度（4）
    output_dim = env.action_space.n  # 动作数量（2）

    # 初始化策略网络和优化器
    policy = PolicyNetwork(input_dim, hidden_dim=128, output_dim=output_dim)
    optimizer = optim.Adam(policy.parameters(), lr=0.01)  # Adam优化器，学习率为0.01

    # 训练循环，运行1000个回合
    for episode in range(1000):
        # 重置环境，开始新回合
        state, _ = env.reset()  # 新版 gym 返回 (state, info)
        # 初始化轨迹存储
        states, actions, rewards = [], [], []
        done = False

        # 采样一条完整轨迹
        while not done:
            # 将当前状态转换为张量
            state_tensor = torch.FloatTensor(state)
            # 通过策略网络获取动作概率分布
            action_probs = policy(state_tensor)
            # 从概率分布中采样一个动作
            action = torch.multinomial(action_probs, 1).item()

            # 执行动作，获取下一状态、奖励和结束标志
            next_state, reward, terminated, truncated, _ = env.step(action)
            # 新版 gym 使用 terminated 和 truncated，done 是两者的结合
            done = terminated or truncated

            # 存储轨迹数据
            states.append(state_tensor)
            actions.append(action)
            rewards.append(reward)

            # 更新当前状态
            state = next_state

        # 计算折扣回报
        returns = compute_returns(rewards)

        # 计算策略梯度损失
        policy_loss = 0
        for t in range(len(rewards)):
            # 获取当前时间步的状态和动作
            state = states[t]
            action = actions[t]
            # 获取当前状态下的动作概率分布
            action_probs = policy(state)
            # 计算所选动作的对数概率
            log_prob = torch.log(action_probs[action])
            # 计算损失：-log π(a_t|s_t) * G_t
            # 负号是因为优化器默认最小化，而我们需要最大化期望回报
            policy_loss -= log_prob * returns[t]

        # 更新策略网络参数
        optimizer.zero_grad()  # 清空之前的梯度
        policy_loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 更新网络参数

        # 打印回合信息
        total_reward = sum(rewards)
        print(f'回合 {episode:4d}, 总奖励: {total_reward:7.2f}')

if __name__ == "__main__":
    train_reinforce()

这个代码实现了一个简单的 REINFORCE 算法，训练一个神经网络策略在 CartPole 环境中保持平衡。

---

第七部分：常见问题与解答

1. 策略梯度与 Q-Learning 的区别？

   • Q-Learning 学习动作价值函数，间接推导策略；策略梯度直接优化策略。
   • 策略梯度适合连续动作空间，Q-Learning 更适合离散动作空间。

2. 为什么策略梯度方差高？

   • 蒙特卡洛采样依赖实际回报，环境随机性和奖励波动导致梯度估计不稳定。

3. 如何选择基线？

   • 状态价值函数 $V(s)$ 是常用的基线，可以用神经网络估计。理想的基线应接近真实价值函数。

4. PPO 和 TRPO 有什么优势？

   • PPO 和 TRPO 通过约束更新步长，防止策略变化过大，提高训练稳定性。PPO 更简单，应用更广泛。

---

第八部分：进一步学习的建议

1. 理论学习：

   • 阅读《Reinforcement Learning: An Introduction》（Sutton & Barto），深入理解强化学习理论。
   • 学习策略梯度的高级算法，如 PPO、SAC，参考原始论文。

2. 实践：

   • 使用 OpenAI Gym 或 Stable-Baselines3 实现 PPO、TRPO 等算法。
   • 尝试在复杂环境（如 MuJoCo、Atari）上训练策略梯度模型。

3. 资源：

   • David Silver 的强化学习课程（YouTube 公开课）。
   • OpenAI Spinning Up 教程，提供代码和理论讲解。

---

总结

策略梯度是强化学习中强大且灵活的方法，通过直接优化策略，适用于离散和连续动作空间。从基础的 REINFORCE 到现代的 PPO、SAC，策略梯度方法在理论和实践上都取得了巨大成功。通过理解其数学推导、算法实现和优化技巧，你可以将其应用于各种复杂任务。