强化学习算法:Policy Gradient
paperreading
本文字数：5.9k 字 | 阅读时长 ≈ 24 min
原创 强化学习

强化学习算法:Policy Gradient

paperreading
本文字数：5.9k 字 | 阅读时长 ≈ 24 min

原创 强化学习

强化学习专项：1. 基本概念 · 2. DQN · 3. Policy Gradient（本文） · 4. PPO · 5. 如何理解强化学习

1. 什么是 Policy Gradient？

Policy Gradient（策略梯度） 是一类基于策略（policy-based）的强化学习方法。DQN 和 Policy Gradient 最终都希望获得更高的期望折扣回报，区别在于实现方式：DQN 先学习 \(Q(s,a)\)，再选择 Q 值最大的动作；Policy Gradient 则直接优化参数化策略 \(\pi_\theta(a\mid s)\)。

本节介绍的 REINFORCE 属于 on-policy 算法：训练数据由当前策略采集，策略更新后需要重新采样。它直接把当前策略的期望折扣回报写成优化目标：

\[ J(\theta) =\mathbb{E}_{\tau\sim\pi_\theta}[U_0],\qquad U_0=\sum_{t=0}^{T-1}\gamma^t r_t \]

下面先把“最大化期望回报”推导成可以计算的策略梯度，再介绍 REINFORCE 怎样采样、计算回报并完成一次更新。

2. 策略梯度的基本思想

上面的 \(J(\theta)\) 说明了最终目标，但它还不是可以直接交给 PyTorch 的 loss。监督学习可以比较 prediction 和 ground truth；策略梯度没有“正确动作”标签，只有策略采样出的动作和环境返回的奖励。因此，需要从 \(J(\theta)\) 出发，推导网络参数应该沿哪个方向更新：

\[ J(\theta) =\mathbb{E}_{\tau\sim\pi_\theta}[U_0] \]

也就是调整策略参数 \(\theta\)，让当前策略产生的轨迹获得更高的平均回报。后面的推导都在解决一个问题：

如何把“最大化期望回报”转换成神经网络能够计算的梯度？

1. 从一步决策理解 \(J(\theta)\)

假设一个状态只有“向左”和“向右”两个动作：

\[ \pi_\theta(\text{向左}\mid s)=0.7,\qquad \pi_\theta(\text{向右}\mid s)=0.3 \]

向左获得奖励 10，向右获得奖励 2，那么当前策略的期望奖励为：

\[ J(\theta) =0.7\times10+0.3\times2 =7.6 \]

网络训练的目标不是拟合某个正确动作，而是调整两个动作的概率，让 \(J(\theta)\) 变大。对它求导：

\[ \nabla_\theta J(\theta) =10\nabla_\theta\pi_\theta(\text{向左}\mid s) + 2\nabla_\theta\pi_\theta(\text{向右}\mid s) \]

奖励更高的动作会产生更大的正向影响。一步决策只有两个结果，可以直接枚举；多步决策会产生大量可能的动作序列，因此接下来需要把同样的思路扩展到完整轨迹。

2. 将期望展开为所有可能的轨迹

当任务包含多个时间步时，一次运行结果不再由单个动作决定，而是由整条轨迹共同决定。一条轨迹记作：

\[ \tau=(s_0,a_0,r_0,s_1,a_1,r_1,\ldots,s_T) \]

这里和 DQN 的区别在于训练时怎样使用轨迹。DQN 与环境交互时也会产生完整轨迹，但会将其拆成单步 transition：

\[ (s_t,a_t,r_t,s_{t+1},d_t) \]

再存入 replay buffer。标准 DQN 每次使用单步数据训练，并通过 \(Q(s_{t+1},a)\) 估计后面的长期价值。REINFORCE 则需要等完整 episode 结束，根据整条轨迹计算每一步的回报 \(U_t\)，所以这里直接从轨迹概率开始推导。

它出现的概率为 \(p_\theta(\tau)\)，整条轨迹的折扣回报为 \(U_0(\tau)\)。将期望展开：

\[ J(\theta) =\sum_\tau p_\theta(\tau)U_0(\tau) \]

这里的两个量含义不同：

例如，假设任务只有三条可能的轨迹：

轨迹 出现概率 \(p_\theta(\tau)\) 回报 \(U_0(\tau)\)
\(\tau_1\) 0.2 10
\(\tau_2\) 0.3 2
\(\tau_3\) 0.5 -1

那么：

\[ J(\theta) =0.2\times10+0.3\times2+0.5\times(-1) =2.1 \]

\(J(\theta)\) 就是所有轨迹回报按照出现概率计算的加权平均。对于一条已经发生的完整轨迹，它的 \(U_0(\tau)\) 是确定的；REINFORCE 不直接训练这个回报，而是训练策略网络参数 \(\theta\)，改变动作概率，进而让高回报轨迹更容易出现。

对参数求导：

\[ \nabla_\theta J(\theta) =\sum_\tau U_0(\tau)\nabla_\theta p_\theta(\tau) \]

问题在于，直接计算所有轨迹及其概率的梯度几乎不可能。实际训练只能让策略跑出有限条轨迹，因此需要把这个求和改写成能够通过采样平均来估计的形式。

3. 怎样把轨迹梯度改写成可以采样的形式？

按照当前策略采样轨迹，本质上是在计算关于 \(p_\theta(\tau)\) 的期望。期望展开后应该具有下面的结构：

\[ \sum_\tau p_\theta(\tau) \left[\cdots\right] \]

但上一节得到的是：

\[ \nabla_\theta J(\theta) =\sum_\tau U_0(\tau)\nabla_\theta p_\theta(\tau) \]

其中没有作为采样概率的 \(p_\theta(\tau)\)。为了把它变成期望，可以乘上一个值为 1 的 \(p_\theta(\tau)/p_\theta(\tau)\)：

\[ \nabla_\theta J(\theta) =\sum_\tau p_\theta(\tau)U_0(\tau) \frac{\nabla_\theta p_\theta(\tau)} {p_\theta(\tau)} \]

这时分式正好等于 log 概率的梯度：

\[ \nabla_\theta\log p_\theta(\tau) =\frac{\nabla_\theta p_\theta(\tau)} {p_\theta(\tau)} \]

代入以后：

\[ \begin{aligned} \nabla_\theta J(\theta) &= \sum_\tau p_\theta(\tau)U_0(\tau) \nabla_\theta\log p_\theta(\tau) \\ &= \mathbb{E}_{\tau\sim\pi_\theta} \left[ U_0(\tau) \nabla_\theta\log p_\theta(\tau) \right] \end{aligned} \]

最后一个等号只是把加权求和重新写成期望。对于按照当前策略产生的离散轨迹，任意函数 \(f(\tau)\) 的期望都可以写成：

\[ \mathbb{E}_{\tau\sim\pi_\theta} \left[f(\tau)\right]= \sum_\tau p_\theta(\tau)f(\tau) \]

\(\tau\sim\pi_\theta\) 表示使用策略 \(\pi_\theta\) 与环境交互，由此产生的轨迹概率就是 \(p_\theta(\tau)\)。这里令：

\[ f(\tau)= U_0(\tau) \nabla_\theta\log p_\theta(\tau) \]

代入期望的定义，就会得到前面那一项求和：每条轨迹上的 \(U_0(\tau)\nabla_\theta\log p_\theta(\tau)\)，按照轨迹出现的概率 \(p_\theta(\tau)\) 加权。因此，这个等号在理论上是精确的。

实际训练无法枚举所有轨迹，只能使用当前策略采样的 \(N\) 条轨迹近似这个期望：

\[ \mathbb{E}_{\tau\sim\pi_\theta} \left[ U_0(\tau)\nabla_\theta\log p_\theta(\tau) \right] \approx \frac{1}{N} \sum_{i=1}^{N} U_0(\tau_i) \nabla_\theta\log p_\theta(\tau_i) \]

上面的改写使用了下面这个关系：

\[ \nabla p=p\nabla\log p \]

这一步通常称为 log-derivative trick。

关键原因：采样平均只能直接估计期望形式

采样平均只能直接估计形如 \(\mathbb{E}_{\tau\sim\pi_\theta}[f(\tau)]\) 的量，也就是 \(\sum_\tau p_\theta(\tau)f(\tau)\)。因此需要利用 log-derivative trick，把原来的梯度改写成带有 \(p_\theta(\tau)\) 的期望形式。

改写以后，代码不需要显式计算完整的 \(p_\theta(\tau)\)：概率较大的轨迹本来就更容易被当前策略采到，采样过程已经完成了概率加权。对于每条采样轨迹，只需要计算回报和策略网络的 log probability，再对样本取平均。

4. 为什么最终只需要对动作概率求导？

上一步仍然包含整条轨迹的概率 \(p_\theta(\tau)\)，但策略网络实际输出的是每个状态下的动作概率 \(\pi_\theta(a_t\mid s_t)\)。为了把公式连接到网络输出，需要先将轨迹概率分解：

\[ p_\theta(\tau) =p(s_0) \prod_{t=0}^{T-1} \pi_\theta(a_t\mid s_t) P(s_{t+1},r_t\mid s_t,a_t) \]

初始状态分布 \(p(s_0)\) 以及环境返回下一状态和奖励的概率 \(P(s_{t+1},r_t\mid s_t,a_t)\) 都由环境决定，不包含策略参数 \(\theta\)。只有动作概率 \(\pi_\theta(a_t\mid s_t)\) 来自策略网络，因此：

\[ \nabla_\theta\log p_\theta(\tau) =\sum_{t=0}^{T-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) \]

于是，完整轨迹回报对应的策略梯度为：

\[ \nabla_\theta J(\theta) =\mathbb{E}_{\tau\sim\pi_\theta} \left[ U_0 \sum_{t=0}^{T-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) \right] \]

这里的期望不是在这一步额外加上的，它来自最初的目标 \(J(\theta)=\mathbb{E}_{\tau\sim\pi_\theta}[U_0]\)。策略采样动作、环境产生下一状态时都可能存在随机性，因此同一个策略可以产生多条不同轨迹。我们要优化的是当前策略反复运行时的平均回报，而不是某一次运行的结果。

具体来说，需要对每条可能轨迹算出括号中的梯度，再按照该轨迹出现的概率求平均。将期望展开，就是：

\[ \nabla_\theta J(\theta) =\sum_\tau p_\theta(\tau) \left[ U_0(\tau) \sum_{t=0}^{T-1} \nabla_\theta\log\pi_\theta(a_t\mid s_t) \right] \]

理论上需要对所有轨迹求这个平均；实际训练无法枚举所有轨迹，所以会用当前策略采样一条或多条轨迹来近似这个期望。

这里的 \(U_0\) 是整条轨迹从起点开始计算的回报。对于一条已经采样出的轨迹，它只有一个值，轨迹中的所有动作暂时都使用同一个 \(U_0\) 评价，因此可以写在时间步求和的外面。展开以后，其实就是每一个动作的梯度都乘上同一个 \(U_0\)。

这一步很重要：我们不需要对环境求导，只需要对策略网络输出的动作概率求导。

5. 为什么每个动作使用自己的 \(U_t\)？

到这里，公式仍然用整条轨迹的 \(U_0\) 评价所有动作。但对时间步 \(t\) 的动作来说，执行该动作之前已经获得的奖励不可能再被它改变。因此，可以去掉过去的奖励，只使用从当前时间步开始的 reward-to-go：

\[ U_t =\sum_{k=t}^{T-1} \gamma^{k-t}r_k \]

\(U_t\) 会随时间步变化：\(a_0\) 使用 \(U_0\)，\(a_1\) 使用 \(U_1\)，依次类推。不同动作对应的权重不再相同，所以 \(U_t\) 不能像前面的 \(U_0\) 一样放在求和外面，而要放进每个时间步自己的梯度项中。

常见的 REINFORCE 梯度估计写成：

核心公式：REINFORCE 策略梯度

\[ \widehat{\nabla_\theta J} =\sum_t U_t \nabla_\theta\log\pi_\theta(a_t\mid s_t) \]

公式上的帽子表示这不再是对所有可能轨迹精确求期望，而是用实际采样到的一条或一批轨迹估计梯度。因此，外层期望在代码中会变成对采样数据计算并取平均。

它表达的含义是：

按照本文对 \(U_t\) 的折扣定义，如果严格对应 \(J(\theta)=\mathbb{E}[U_0]\)，第 \(t\) 项还会带有外层权重 \(\gamma^t\)。许多简化实现，包括下面的 CartPole 代码，会省略这个外层权重，只保留 \(U_t\) 内部的折扣。这会改变不同时间步的权重，但不会改变“用后续回报调整实际动作概率”的核心逻辑。

6. 如何从策略梯度得到代码中的 loss？

前面的核心公式给出了增大 \(J(\theta)\) 所需的梯度方向，但 PyTorch 优化器需要一个标量 loss，并且默认执行梯度下降。因此，对于一批采样数据，可以构造下面的 surrogate loss，也就是用于产生所需梯度的替代损失：

训练公式：代码中实际计算的策略损失

\[ L(\theta) =-\frac{1}{N} \sum_t U_t\log\pi_\theta(a_t\mid s_t) \]

这里的 \(N\) 是当前 batch 中的时间步总数；如果 batch 包含多条轨迹，就把这些轨迹中的时间步合在一起取平均。

对它求导：

\[ \nabla_\theta L(\theta) =-\frac{1}{N} \sum_t U_t \nabla_\theta\log\pi_\theta(a_t\mid s_t) \]

梯度下降执行：

\[ \theta \leftarrow \theta-\alpha\nabla_\theta L(\theta) \]

代入负号后，相当于沿着增大 \(J(\theta)\) 的方向更新参数。因此，这个 loss 不是拿 prediction 和 ground truth 比较，也不是将 \(J(\theta)\) 完全等价地改写成另一个数值，而是：

构造一个可供自动求导的目标，使它的梯度产生我们需要的策略梯度方向。

代码中两部分数据的来源也不同：

对应到 PyTorch：

action_probs = policy_network(states)
log_probs = Bernoulli(action_probs).log_prob(actions)
policy_loss = -(log_probs * normalized_returns).mean()

log_probs 是网络中可求导的部分，normalized_returns 是由采样奖励计算出的权重。执行 policy_loss.backward() 后，梯度会提高高回报动作的概率，并降低相对低回报动作的概率。

整个过程可以概括为：

定义目标：最大化期望回报 J(θ)
              ↓
对 J(θ) 求导，并使用 log trick 改写
              ↓
网络计算 log πθ(a_t | s_t)
环境采样 s_t、a_t、r_t 和 U_t
              ↓
构造 surrogate loss
              ↓
loss.backward()

3. REINFORCE 算法

第二章已经把期望回报推导成了可以通过采样计算的 loss。REINFORCE 就是按照这套结果训练：收集完整 episode，用后续回报评价每个实际动作，再更新策略网络。

训练流程如下：

  1. 初始化策略网络 \(\pi_\theta(a\mid s)\)；
  2. 使用当前策略采样一条或多条完整轨迹；
  3. 从后向前计算每个时间步的回报 \(U_t\)；
  4. 计算策略损失 \(L=-\mathbb{E}[\log\pi_\theta(a_t\mid s_t)U_t]\)；
  5. 反向传播并更新参数 \(\theta\)；
  6. 丢弃旧轨迹，使用更新后的策略重新采样。

下图对应这六步。策略网络先输出动作概率，再从中采样实际动作；环境返回的 reward 会随轨迹一起保存。episode 结束后计算每一步的 \(U_t\)，再用它评价当时实际执行的动作，组成策略损失并更新参数。

实际训练中常对一批轨迹的 \(U_t\) 做标准化。标准化后的值大于 0，表示该时间步的回报高于当前 batch 的平均水平；小于 0，则表示低于平均水平。这样可以降低梯度方差，但此时用于更新的是相对回报，而不是原始 \(U_t\) 的正负。

3.1 episode 长度为什么会影响 \(U_t\)？

在 CartPole 中，每存活一步都会获得 reward 1。一个 episode 在时间步 \(T\) 结束时：

\[ U_t =1+\gamma+\gamma^2+\cdots+\gamma^{T-t-1} =\frac{1-\gamma^{T-t}}{1-\gamma} \]

因此，episode 越长，前面时间步的 \(U_t\) 通常越大。这不是计算错误，而是 CartPole 的任务目标：杆保持得越久，得到的回报越高。

同一个 episode 中，越靠后的 \(U_t\) 又会越小，因为后面剩余的 reward 更少。这只表示“从当前位置开始还能观察到多少回报”，并不一定表示后面的动作本身更差。

不同长度的 episode 放在一起标准化时，平均值确实会受到长 episode 影响。假设 \(\gamma=1\)，两个 episode 的回报分别为：

\[ \begin{aligned} \text{episode A}:&\quad [4,3,2,1] \\ \text{episode B}:&\quad [2,1] \end{aligned} \]

全部 \(U_t\) 的平均值为：

\[ \mu_U =\frac{4+3+2+1+2+1}{6} =2.17 \]

此时，长 episode 前面的回报高于平均值，而短 episode 和靠近结束位置的回报大多低于平均值。长 episode 还会产生更多时间步，因此在 batch 中占有更多样本。使用整个 batch 的均值作为基准能够降低梯度方差，但它没有区分状态和时间位置，只是一种比较简单的处理方式。

标准化时，减去均值是为了判断每个 \(U_t\) 高于还是低于 batch 的平均水平，除以标准差只是为了控制训练信号的数值尺度。这两步都不会让某一次采样得到的 \(U_t\) 本身变得更准确。

episode 只有一两步也不一定表示回报计算不准确：

一条轨迹的 \(U_t\) 是这次运行的实际结果，但它只是期望回报的一次随机样本。增加 episode 数量能够让整体估计更稳定，并不代表 episode 本身越长，估计就一定越准确。

这也是 Actor-Critic 引入状态价值的原因。它不再统一使用整个 batch 的平均回报作为基准，而是针对当前状态计算：

\[ A_t =U_t-V(s_t) \]

\(V(s_t)\) 表示当前状态下通常能够获得多少回报，因此 \(A_t\) 更接近“这次动作相对于当前状态的正常表现是好还是差”。PPO 会在此基础上进一步构造优势估计。

3.2 用一条轨迹计算 loss

假设一条轨迹包含 3 个时间步，每一步都选择“向右”，奖励为 \((0,0,1)\)，并令 \(\gamma=1\)：

\(t\) 动作 \(a_t\) 奖励 \(r_t\) 回报 \(U_t\) 动作概率 \(\pi_\theta(a_t\mid s_t)\) \(\log\pi_\theta(a_t\mid s_t)\)
0 向右 0 1 0.5 -0.693
1 向右 0 1 0.6 -0.511
2 向右 1 1 0.7 -0.357

代入策略损失：

\[ L=-\left[(-0.693)\times1+(-0.511)\times1+(-0.357)\times1\right] =1.561 \]

这里使用的是原始回报，三个 \(U_t\) 都为正，因此梯度下降会提高这三个已选动作的概率。下面的 CartPole 代码会进一步标准化回报，使策略根据每个时间步相对于 batch 平均水平的表现进行更新。

3.3 REINFORCE 的特点

Actor-Critic 和 PPO 会进一步引入价值函数、优势函数等机制，以降低方差并提高训练稳定性。

4. CartPole 代码实现

前面用一条短轨迹手动计算了 loss，下面将相同过程扩展到完整训练。代码与公式的对应关系为：action_probs 对应 \(\pi_\theta(a_t\mid s_t)\)，log_probs 对应 \(\log\pi_\theta(a_t\mid s_t)\)，normalized_returns 对应标准化后的 \(U_t\)，policy_loss 对应策略损失 \(L\)。

4.1 一次实际训练是怎样执行的？

下面固定随机种子运行一次代码。数值用于展示计算过程；实际训练中的动作和 episode 长度会随采样改变。

运行配置 数值
随机种子 7
折扣因子 \(\gamma\) 0.99
学习率 0.01
每次更新收集的 episode 5
策略网络参数量 1,057

策略参数 \(\theta\) 在收集完 5 个完整 episode 后更新一次。下面用动作概率观察这次更新产生的效果。

1. 策略网络根据概率采样动作

初始状态和第一次交互的数据为：

量 数值
\(s_0\) \([0.0125,\ 0.0397,\ 0.0276,\ -0.0275]\)
\(\pi_\theta(a=0\mid s_0)\) 0.5555
\(\pi_\theta(a=1\mid s_0)\) 0.4445
采样动作、reward \(a_0=0,\ r_0=1\)
\(s_1\) \([0.0133,\ -0.1558,\ 0.0270,\ 0.2738]\)

动作按照概率采样，并不是直接选择概率最大的动作。代码会继续保存状态、动作和 reward，但必须等 episode 结束后才能计算每一步的 \(U_t\)，此时策略网络还不会更新。

2. episode 结束后，从后向前计算 \(U_t\)

第一个 episode 在第 24 步结束，奖励序列为 24 个 1。代码从最后一步开始反向计算：

\[ \begin{aligned} U_{23}&=1 \\ U_{22}&=1+0.99\times U_{23}=1.99 \\ U_{21}&=1+0.99\times U_{22}=2.9701 \\ &\ \vdots \\ U_0&=21.4322 \end{aligned} \]

第一个 episode 共 24 步，所以 \(U_{23}=1\)，而 \(U_0=21.4322\)。越靠前的时间步包含的未来奖励通常越多。

3. 收集 5 个完整 episode

环境在每个 episode 结束后重置，但策略保持不变，直到收集满 5 个 episode：

episode 时间步数 \(U_0\) batch 累计时间步
1 24 21.4322 24
2 12 11.3615 36
3 54 41.8834 90
4 13 12.2479 103
5 47 37.6475 150

这批数据共包含 150 个时间步，并且全部来自更新前的同一个策略。

4. 标准化回报并计算 loss

这批回报的统计量为：

时间步数 \(\mu_U\) \(\sigma_U\)
150 17.7822 11.6999

每个回报转换为：

\[ \widetilde{U}_t =\frac{U_t-\mu_U}{\sigma_U} \]

其中两个时间步形成 loss 的过程如下：

时间步 实际动作 该动作的概率 \(U_t\) \(\widetilde{U}_t\) \(\log\pi_\theta(a_t\mid s_t)\) loss 中的这一项
3 1 0.4548 19.0272 0.1064 -0.7878 0.0838
23 0 0.5498 1.0000 -1.4344 -0.5982 -0.8581

最后一列计算的是：

\[ -\log\pi_\theta(a_t\mid s_t)\widetilde{U}_t \]

标准化回报为正时提高实际动作的概率，为负时降低实际动作的概率。代码对全部 150 项取平均，本次运行得到：

\[ L=0.0106 \]

这个 loss 用于产生策略梯度，不能直接当作 episode 得分。

5. 更新策略，然后重新采样

执行一次 optimizer.step() 后，两个示例状态中实际动作的概率变化如下：

时间步 实际动作 更新前概率 更新后概率
3 1 0.4548 0.6782
23 0 0.5498 0.2885

概率变化符合标准化回报给出的方向，但网络参数实际由 150 项的平均 loss 共同更新。更新后丢弃这批数据，使用新策略重新采样：

采集 5 个 episode → 计算并标准化 U_t → 计算平均 loss
→ 更新一次策略 → 丢弃旧数据并重新采样

4.2 完整代码

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
from torch.distributions import Bernoulli


class PolicyNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(4, 24), nn.ReLU(),
            nn.Linear(24, 36), nn.ReLU(),
            nn.Linear(36, 1), nn.Sigmoid(),
        )

    def forward(self, x):
        return self.network(x).squeeze(-1)


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    env = gym.make("CartPole-v1")

    random_seed = 7
    np.random.seed(random_seed)
    torch.manual_seed(random_seed)
    env.action_space.seed(random_seed)

    num_episodes = 5000
    episodes_per_update = 5
    learning_rate = 0.01
    gamma = 0.99

    policy_network = PolicyNetwork().to(device)
    optimizer = torch.optim.RMSprop(
        policy_network.parameters(), lr=learning_rate
    )

    num_updates = num_episodes // episodes_per_update
    for update in range(num_updates):
        batch_states = []
        batch_actions = []
        batch_returns = []
        episode_scores = []

        # 使用同一个策略收集若干完整 episode
        for episode_index in range(episodes_per_update):
            first_episode = update == 0 and episode_index == 0
            state, _ = env.reset(
                seed=random_seed if first_episode else None
            )
            episode_states = []
            episode_actions = []
            episode_rewards = []
            done = False

            while not done:
                state_tensor = torch.as_tensor(
                    state, dtype=torch.float32, device=device
                )
                with torch.no_grad():
                    action_prob = policy_network(state_tensor)
                    action = int(Bernoulli(action_prob).sample().item())

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

                episode_states.append(state)
                episode_actions.append(float(action))
                episode_rewards.append(reward)
                state = next_state

            # 从后向前计算当前 episode 的 U_t
            episode_returns = []
            discounted_return = 0.0
            for reward in reversed(episode_rewards):
                discounted_return = reward + gamma * discounted_return
                episode_returns.insert(0, discounted_return)

            batch_states.extend(episode_states)
            batch_actions.extend(episode_actions)
            batch_returns.extend(episode_returns)
            episode_scores.append(sum(episode_rewards))

        states = torch.as_tensor(
            np.asarray(batch_states), dtype=torch.float32, device=device
        )
        actions = torch.as_tensor(
            batch_actions, dtype=torch.float32, device=device
        )
        returns = torch.as_tensor(
            batch_returns, dtype=torch.float32, device=device
        )
        normalized_returns = (
            returns - returns.mean()
        ) / (returns.std(unbiased=False) + 1e-8)

        action_probs = policy_network(states)
        log_probs = Bernoulli(action_probs).log_prob(actions)
        policy_loss = -(log_probs * normalized_returns).mean()

        optimizer.zero_grad()
        policy_loss.backward()
        optimizer.step()

        if update % 10 == 0:
            mean_score = sum(episode_scores) / episodes_per_update
            trained_episodes = (update + 1) * episodes_per_update
            print(
                f"Episode {trained_episodes:4d} | "
                f"Mean score {mean_score:6.1f} | "
                f"Loss {policy_loss.item():7.4f}"
            )

    env.close()


if __name__ == "__main__":
    main()

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: Oct 04, 2026 05:40 pm

← → ↓ ↑
Codex 的输入 Token 都花在哪了？
Oct 25, 2026
utils
BM25：从词频、IDF 到 RAG 关键词检索
Sep 08, 2026
statistic
Agent 工程：Prompt、Context、Harness 与 Loop
Sep 06, 2026
utils
InfoNCE Loss
Sep 06, 2026
pytorch
从零开始学习 Codex Skill
Aug 15, 2026
utils
向量表征与逻辑回归的分类原理
Aug 13, 2026
statistic