策略梯度算法（Policy Gradient）理论基础及REINFORCE算法代码实现

Q-learning、DQN及DQN改进算法都是基于价值（value-based）的方法，其中Q-learning是处理有限状态的算法，而DQN可以用来解决连续状态的问题。在强化学习中，除了基于值函数的方法，还有一支非常经典的方法，那就是基于策略（policy-based）的方法。对比两者，基于值函数的方法主要是学习值函数，然后根据值函数导出一个策略，学习过程中并不存在一个显式的策略；而基于策略的方法则是直接显式地学习一个目标策略。策略梯度是基于策略的方法的基础。

在学习这个算法之前，我们先来解决如下两个问题。

为什么要用基于策略的学习？

• 基于策略的学习可能会具有更好的收敛性，这是因为基于策略的学习虽然每次只改善一点点，但总是朝着好的方向在改善；但是基于值函数的方法在后期会一直围绕最优价值函数持续小的震荡而不收敛。
• 在对于那些拥有高维度或连续状态空间来说，使用基于价值函数的学习在得到价值函数后，制定策略时，需要比较各种行为对应的价值大小，这样如果行为空间维度较高或者是连续的，则从中比较得出一个有最大价值函数的行为这个过程就比较难了，这时候使用基于策略的学习就高效的多。
• 能够学到一些随机策略，但是基于价值函数的学习通常是学不到随机策略的。
• 有时候计算价值函数很困难。比如当小球从空中掉下来你需要通过左右移动去接住它时，计算小球在某一个位置(状态)时采取什么样的动作是很困难的。但是基于策略函数就简单了，只需要朝着小球落地的方向移动修改策略就好了。

什么时候使用基于价值的学习？什么时候使用基于策略的学习？

这个问题当然要具体问题具体分析了，我们必须要根据需要评估的问题的特点来决定使用哪一种学习方式。随机策略有时是最优策略。比如剪刀石头布这个游戏，如果你是按照某一种策略来出拳的话，很容易让别人抓住你的规律，然后你就会输了。所以最好的策略就是随机出拳，让别人猜不到。
这里要分清确定性策略和随机性策略：

• 所谓的确定性策略，是说只要给定一个状态 s s s，就会输出一个具体的动作 a a a，而且无论什么时候到达状态 s s s，输出的动作 a a a都是一样的。
• 而随机策略是指，给定一个状态 s s s，输出在这个状态下可以执行的动作的概率分布。即使在相同状态下，每次采取的动作也很可能是不一样的。

了解了这些之后，正式开始今天的主题。下面我们对策略梯度（Policy Gradient）算法进行推导。

不管什么类型的方法，强化学习的最终目的都是要使得得到的奖励最大化，因此假设这个目标函数为 J ( θ ) J(theta) J(θ)，那么最终的目的就是为了最大化这个目标函数，将轨迹的期望回报展开，可以得到：
J ( θ ) = E τ ∼ π ( θ ) [ r ( τ ) ] = ∫ τ ∼ π ( θ ) π θ ( τ ) r ( τ ) d τ J(theta)=E_{tausimpi(theta)}[r(tau)]=int_{tausimpi(theta)} pi_{theta}(tau) r(tau) d tau J(θ)=Eτ∼π(θ)​[r(τ)]=∫τ∼π(θ)​πθ​(τ)r(τ)dτ
下面对公式求导，因为积分和求导运算可以互换
∇ θ J ( θ ) = ∇ θ ∫ τ ∼ π ( θ ) π θ ( τ ) r ( τ ) d τ = ∫ τ ∼ π ( θ ) ∇ θ π θ ( τ ) r ( τ ) d τ nabla_{theta} J(theta)=nabla_{theta} int_{tausimpi(theta)} pi_{theta}(tau) r(tau) d tau=int_{tausimpi(theta)} nabla_{theta} pi_{theta}(tau) r(tau) d tau ∇θ​J(θ)=∇θ​∫τ∼π(θ)​πθ​(τ)r(τ)dτ=∫τ∼π(θ)​∇θ​πθ​(τ)r(τ)dτ
因为积分的缘故，这个形式不方便直接计算，可以对其做一个变换，这里可以用到对数求导的基本公式：
∇ x log ⁡ y = 1 y ∇ x y nabla_{x}log y=frac{1}{y}nabla_{x}y ∇x​logy=y1​∇x​y
经过变换可以得到：
y ∇ x log ⁡ y = ∇ x y ynabla_{x}log y=nabla_{x}y y∇x​logy=∇x​y
故有：
∇ θ π θ ( τ ) = π θ ( τ ) ∇ θ log ⁡ π θ ( τ ) nabla_{theta}pi_{theta}(tau)=pi_{theta}(tau)nabla_{theta}logpi_{theta}(tau) ∇θ​πθ​(τ)=πθ​(τ)∇θ​logπθ​(τ)
带入前面的公式，有：
∇ θ J ( θ ) = ∫ τ ∼ π ( θ ) ∇ θ π θ ( τ ) r ( τ ) d τ = ∫ τ ∼ π ( θ ) π θ ( τ ) ∇ θ log ⁡ π θ ( τ ) r ( τ ) d τ begin{aligned}nabla_{theta}J(theta)&=int_{tausimpi(theta)}nabla_{theta}pi_{theta}(tau)r(tau)dtau\&=int_{tausimpi(theta)}pi_{theta}(tau)nabla_{theta}logpi_{theta}(tau)r(tau)dtauend{aligned} ∇θ​J(θ)​=∫τ∼π(θ)​∇θ​πθ​(τ)r(τ)dτ=∫τ∼π(θ)​πθ​(τ)∇θ​logπθ​(τ)r(τ)dτ​
将轨迹 τ tau τ展开，可以得到：
π θ ( τ ) = π ( s 0 , a 0 , … , s T , a T ) = p ( s 0 ) ∏ t = 0 T π θ ( a t ∣ s t ) p ( s t + 1 ∣ s t , a t ) pi_{theta}(tau)=pileft(s_{0},a_{0},ldots,s_{T},a_{T}right)=pleft(s_{0}right)prod_{t=0}^{T}pi_{theta}left(a_{t}|s_{t}right)pleft(s_{t+1}|s_{t},a_{t}right) πθ​(τ)=π(s0​,a0​,…,sT​,aT​)=p(s0​)t=0∏T​πθ​(at​∣st​)p(st+1​∣st​,at​)
所以，
∇ θ log ⁡ [ π ( τ ) ] = ∇ θ log ⁡ [ p ( s 0 ) ∏ t = 0 T π θ ( a t ∣ s t ) p ( s t + 1 ∣ s t , a t ) ] = ∇ θ [ log ⁡ p ( s 0 ) + ∑ t = 0 T log ⁡ π θ ( a t ∣ s t ) + ∑ t = 0 T log ⁡ p ( s t + 1 ∣ s t , a t ) ] = ∑ t = 0 T ∇ θ log ⁡ π θ ( a t ∣ s t ) begin{aligned}nabla_{theta}log[pi(tau)]&=nabla_{theta}logleft[pleft(s_{0}right)prod_{t=0}^{T}pi_{theta}left(a_{t}|s_{t}right)pleft(s_{t+1}|s_{t},a_{t}right)right]\&=nabla_{theta}left[log pleft(s_{0}right)+sum_{t=0}^{T}logpi_{theta}left(a_{t}|s_{t}right)+sum_{t=0}^{T}log pleft(s_{t+1}|s_{t},a_{t}right)right]\&=sum_{t=0}^{T}nabla_{theta}logpi_{theta}left(a_{t}|s_{t}right)end{aligned} ∇θ​log[π(τ)]​=∇θ​log[p(s0​)t=0∏T​πθ​(at​∣st​)p(st+1​∣st​,at​)]=∇θ​[logp(s0​)+t=0∑T​logπθ​(at​∣st​)+t=0∑T​logp(st+1​∣st​,at​)]=t=0∑T​∇θ​logπθ​(at​∣st​)​

最后一步是因为第一项和第三项与 θ theta θ无关。

最后，再使用蒙特卡罗法，将公式中的期望用蒙特卡罗近似的方式进行替换，得到求解梯度的最终形式：
∇ θ J ( θ ) = ∫ τ ∼ π ( θ ) π θ ( τ ) ∇ θ log ⁡ π θ ( τ ) r ( τ ) d τ = E τ ∼ π θ ( τ ) [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a i , t ∣ s i , t ) ∑ t = 0 T r ( s i , t , a i , t ) ] = 1 N ∑ i = 1 N [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a i , t ∣ s i , t ) ∑ t = 0 T r ( s i , t , a i , t ) ] begin{aligned}nabla_{theta}J(theta)&=int_{tausimpi(theta)}pi_{theta}(tau)nabla_{theta}logpi_{theta}(tau)r(tau)dtau\&=E_{tausimpi_{theta}(tau)}left[sum_{t=0}^{T}nabla_{theta}logpi_{theta}left(a_{i,t}|s_{i,t}right)sum_{t=0}^{T}rleft(s_{i,t},a_{i,t}right)right]\&=frac{1}{N}sum_{i=1}^{N}left[sum_{t=0}^{T}nabla_{theta}logpi_{theta}left(a_{i,t}|s_{i,t}right)sum_{t=0}^{T}rleft(s_{i,t},a_{i,t}right)right]end{aligned} ∇θ​J(θ)​=∫τ∼π(θ)​πθ​(τ)∇θ​logπθ​(τ)r(τ)dτ=Eτ∼πθ​(τ)​[t=0∑T​∇θ​logπθ​(ai,t​∣si,t​)t=0∑T​r(si,t​,ai,t​)]=N1​i=1∑N​[t=0∑T​∇θ​logπθ​(ai,t​∣si,t​)t=0∑T​r(si,t​,ai,t​)]​
这就完成了对梯度的求解，然后就是用梯度下降法对参数进行更新。

但是对于上式，由于这个最后一项的加权项的存在，会使得策略梯度的方差特别大。不论哪个时间段，我们都要用策略的梯度乘以后面这个所有时刻的回报值总和，这样做显然不合理，所以我们利用到当前的决策不能影响之前的回报的原理： t t t时刻我们完成决策之后，它最多只能影响 t t t时刻之后的回报，不会影响之前的回报，所以我们不应该将之前的回报和计算在梯度中，公式改写为：
∇ θ J ( θ ) = 1 N ∑ i = 1 N [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a i , t ∣ s i , t ) ( ∑ t ′ = t T r ( s i , t ′ , a i , t ′ ) ) ] nabla_{theta}J(theta)=frac{1}{N}sum_{i=1}^{N}left[sum_{t=0}^{T}nabla_{theta}logpi_{theta}left(a_{i,t}|s_{i,t}right)left(sum_{t^{prime}=t}^{T}rleft(s_{i,t^{prime}},a_{i,t^{prime}}right)right)right] ∇θ​J(θ)=N1​i=1∑N​[t=0∑T​∇θ​logπθ​(ai,t​∣si,t​)(t′=t∑T​r(si,t′​,ai,t′​))]

从这里可以看出来，策略梯度方法更像是加权版的最大似然优化法。“权重”将直接影响梯度的更新量，这样就会带来以下两个问题：

• 如果计算得出的序列回报数值较大，那么对应的参数更新量就会比较大，这样优化就可能出现一定波动，这些波动很可能影响优化效果；
• 在一些问题中，环境给予的回报始终为正， 那么不论决策如何， 最终累积的长期回报值都是一个正数。换句话说，我们会提升所有的策略，只是对于实际效果并不好的策略，我们为其提升的幅度会有所降低。而初衷是提高能最大化长期回报策略的概率，降低无法最大化长期回报策略的概率。

回到强化学习的目标：提高能最大化长期回报策略的概率，降低无法最大化长期回报策略的概率。将上面的思想转化成策略梯度问题的表述形式，就会变成：让能够最大化长期回报策略的“权重”为正且尽可能的大，让不能最大化长期回报策略的“权重”为负且尽可能地小。
为了实现这个目标，我们可以调整权重的数值和范围，一个简单的方法就是给所有时刻的长期累积回报减去一个偏移量， 这个偏移量也被称为Baseline ，用变量 b b b表示，于是公式就变为：
∇ θ J ( θ ) = 1 N ∑ i = 1 N [ ∑ t = 0 T ∇ θ log ⁡ π θ ( a i , t ∣ s i , t ) ( ∑ t ′ = t T r ( s i , t ′ , a i , t ′ ) − b i , t ′ ) ] nabla_{theta}J(theta)=frac{1}{N}sum_{i=1}^{N}left[sum_{t=0}^{T}nabla_{theta}logpi_{theta}left(a_{i,t}|s_{i,t}right)left(sum_{t^{prime}=t}^{T}rleft(s_{i,t^{prime}},a_{i,t^{prime}}right)-b_{i,t^{prime}}right)right] ∇θ​J(θ)=N1​i=1∑N​[t=0∑T​∇θ​logπθ​(ai,t​∣si,t​)(t′=t∑T​r(si,t′​,ai,t′​)−bi,t′​)]
这个变量可以设计为同一起点地不同序列在同一时刻地长期回报均值，他的公式形式如下：
b i , t ′ = 1 N ∑ i = 1 N ∑ t ′ = t T r ( s i , t ′ , a i , t ′ ) b_{i,t^{prime}}=frac{1}{N}sum_{i=1}^{N}sum_{t^{prime}=t}^{T}rleft(s_{i,t^{prime}},a_{i,t^{prime}}right) bi,t′​=N1​i=1∑N​t′=t∑T​r(si,t′​,ai,t′​)
这样，所有时刻的权重均值变为0 ，就会存在权重为正或为负的行动，同时权重的绝对值也得到了一定的缩小。这相当于对长期回报值期望规零化，对算法的稳定性有一定的帮助。

事实上，引入偏移量并不会使原来的计算有偏，即：
E [ ∇ θ log ⁡ π θ ( τ ) b ] = ∫ τ ∼ π θ ( τ ) π θ ( τ ) ∇ θ log ⁡ π θ ( τ ) b   d τ = ∫ τ ∼ π θ ( τ ) ∇ θ π θ ( τ ) b   d τ = b ∫ τ ∼ π θ ( τ ) ∇ θ π θ ( τ ) d τ = b ∇ θ ∫ τ ∼ π θ ( τ ) π θ ( τ ) d τ = b ∇ θ 1 = 0 begin{aligned} Eleft[nabla_{theta} log pi_{theta}(tau) bright] &=int_{tau sim pi_{theta}(tau)} pi_{theta}(tau) nabla_{theta} log pi_{theta}(tau) b mathrm{~d} tau \ &=int_{tau sim pi_{theta}(tau)} nabla_{theta} pi_{theta}(tau) b mathrm{~d} tau \ &=b int_{tau sim pi_{theta}(tau)} nabla_{theta} pi_{theta}(tau) mathrm{d} tau \ &=b nabla_{theta} int_{tau sim pi_{theta}(tau)} pi_{theta}(tau) mathrm{d} tau \ &=b nabla_{theta} 1 \ &=0 end{aligned} E[∇θ​logπθ​(τ)b]​=∫τ∼πθ​(τ)​πθ​(τ)∇θ​logπθ​(τ)b dτ=∫τ∼πθ​(τ)​∇θ​πθ​(τ)b dτ=b∫τ∼πθ​(τ)​∇θ​πθ​(τ)dτ=b∇θ​∫τ∼πθ​(τ)​πθ​(τ)dτ=b∇θ​1=0​
所以它可以在不影响期望值的同时降低算法的波动性。

Policy Gradient算法的优点：

• 具有很好的收敛性
• 对于高维空间或者是连续空间更加的有效
• 能够对随机策略进行学习

Policy Gradient算法的缺点：

• 很容易在局部最优解上面收敛而得不到全局最优
• 对策略的估计通常具有很大的方差，求解的过程较低效

事实上，现在基本没多少人在用最原始的PG算法了，大多用的都是Actor-Critic家族的算法。

先介绍两种不同的更新方法：

• 蒙特卡洛（Monte-Carlo），属于回合更新。当算法完成一个回合后，每个时刻的奖励 r t r_t rt​都可以获取到，这样就可以计算未来总奖励 G t G_t Gt​。完成一次回合才learn一次。
• 时序差分（Temporal-Difference），属于单步更新，比如Q-learning。每个步骤都learn一下。

我们在上面提到用蒙特卡洛来估计期望，这其实就是REINFORCE算法的思想。REINFORCE的算法流程如下

1. 初始化策略 π pi π的参数 θ theta θ
2. for 每一个完整的episode， { s 1 , a 1 , r 2 , . . . , s T − 1 , a T − 1 , r T } ∼ π θ left{ {{s_1},{a_1},{r_2},...,{s_{T - 1}},{a_{T - 1}},{r_T}} right} sim {pi _theta } {s1​,a1​,r2​,...,sT−1​,aT−1​,rT​}∼πθ​：
   1. for t = 1 t=1 t=1到 t = T − 1 t=T-1 t=T−1：
      • 用蒙特卡洛求经验平均值来计算 v t v_t vt​
      • 更新参数 θ theta θ： θ ← θ + α ∇ θ log ⁡ π θ ( s t , a t ) v t theta leftarrow theta + alpha {nabla _theta }log {pi _theta }left( {{s_t},{a_t}} right){v_t} θ←θ+α∇θ​logπθ​(st​,at​)vt​
   2. 依据更新的 θ theta θ值，按照新的策略 π θ pi_theta πθ​生成新的完整的episode，回到步骤2.
3. 重复以上步骤，从许多个episode中不断更新 θ theta θ，从而得到最优策略 π pi π

在网络的输出层用Softmax函数作用，使得神经网络输出每个动作对应的概率。

class PolicyNet(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.fc_layer = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
        )

    def forward(self, x):
        return F.softmax(self.fc_layer(x), dim=1)

REINFORCE 算法的代码实现，其实重点就在select_action()函数和update()函数，其他的基本差不多。像select_action()函数中的实现，后续很多算法都是这样做的。

class REINFORCE:
    def __init__(self, hidden_dim=128, learning_rate=1e-3, gamma=0.98):
        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

        self.env_name = "CartPole-v0"
        self.env = gym.make(self.env_name)
        state_dim = self.env.observation_space.shape[0]
        action_dim = self.env.action_space.n

        self.env.seed(0)
        torch.manual_seed(0)

        self.policy_net = PolicyNet(state_dim, hidden_dim, action_dim).to(self.device)
        self.optimizer = torch.optim.Adam(self.policy_net.parameters(), lr=learning_rate)  # 使用Adam优化器
        self.gamma = gamma  # 折扣因子
        self.num_episodes = 1000  # 训练的总回合数

    def select_action(self, state):  # 根据动作概率分布随机采样
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        probs = self.policy_net(state)
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        return action.item()

    def update(self, transition_dict):
        reward_list = transition_dict["rewards"]
        state_list = transition_dict["states"]
        action_list = transition_dict["actions"]

        G = 0
        self.optimizer.zero_grad()
        for i in reversed(range(len(reward_list))):  # 从最后一步算起，反向计算
            reward = reward_list[i]
            state = torch.tensor([state_list[i]], dtype=torch.float).to(self.device)
            action = torch.tensor([action_list[i]]).view(-1, 1).to(self.device)

            log_prob = torch.log(self.policy_net(state).gather(1, action))
            G = self.gamma * G + reward  # 每一步的损失函数
            loss = - log_prob * G
            loss.backward()
        self.optimizer.step()

    def run(self, ):
        return_list = []
        for i in range(10):
            with tqdm(total=self.num_episodes // 10, desc=f"Iteration {i}") as pbar:
                for ep in range(self.num_episodes // 10):
                    ep_return = 0
                    transition_dict = {"states": [], "actions": [], "next_states": [], "rewards": [], "dones": []}
                    state = self.env.reset()
                    done = False
                    while not done:
                        action = self.select_action(state)
                        next_state, reward, done, _ = self.env.step(action)
                        transition_dict["states"].append(state)
                        transition_dict["actions"].append(action)
                        transition_dict["next_states"].append(next_state)
                        transition_dict["rewards"].append(reward)
                        transition_dict["dones"].append(done)
                        state = next_state
                        ep_return += reward
                    return_list.append(ep_return)
                    self.update(transition_dict)

                    if (ep + 1) % 10 == 0:
                        pbar.set_postfix({
                            'episode': '%d' % (self.num_episodes / 10 * i + ep + 1),
                            'return': '%.3f' % np.mean(return_list[-10:])
                        })
                    pbar.update(1)

        self.plot(return_list)

    def plot(self, return_list):
        episodes_list = list(range(len(return_list)))
        plt.plot(episodes_list, return_list)
        plt.xlabel('Episodes')
        plt.ylabel('Returns')
        plt.title('REINFORCE on {}'.format(self.env_name))
        plt.show()

代码运行结果如下：

可以发现REINFORCE 算法不是很稳定，这也是它的一个非常大的缺点！后续的Actor-Critic系列的算法对这个算法进行了改进。

参考：

• https://blog.csdn.net/liweibin1994/article/details/79348000
• https://zhuanlan.zhihu.com/p/450690041
• https://blog.csdn.net/sherlocksy/article/details/119490320
• 《动手学强化学习》