利用PPO算法训练超级玛丽AI模型

本文使用PPO增强学习策略，训练一个自动玩超级玛丽的AI模型。在进入主题之前，先介绍几个代码中会用到的python库。

• gym：openai开发的python库，提供了开发和比较增强学习算法的标准API。简单地说，这个库是一个游戏模拟器，可以通过API控制游戏如何进行。
• pytorch: 使用GPU加速的动态神经网络训练程序库。

第一节 PPO 算法

PPO，全称Proximal Policy Optimization Algorithms，是openai提出的增强学习策略，出自论文Proximal Policy Optimization Algorithms，应用于Instruct GPT对齐训练阶段，也是许多大模型的对齐训练策略之一。基本原理是：先利用当前模型计算出一系列动作（或基于随机采样），得到一个状态和动作序列，随后通过离线迭代，重复利用这个序列更新模型。更新的策略为：如果某个动作产生正向激励，那么离线模型将提升在此状态下该动作的概率，反之抑制其概率。这个策略与Q-Learning类似，一般来说，采样状态和动作序列的成本较高，利用同一序列进行多次更新则更有效的利用采样数据。与Q-Learning不同的是，Q-Learning每次计算策略梯度时，仅使用一步迭代后的状态，并以Q矩阵估计长期激励，而PPO则需要利用整个序列，在此基础上计算动作的相对熵、价值误差与策略分布自身的熵三部分之和。同时，为了保证离线模型在更新策略梯度时的平稳性，使用一个小的误差窗口将策略误差控制在一个小范围内。
分别通过PPO和Q-Learning对超级玛丽奥游戏的学习发现，PPO算法的学习效率更高，通过5次迭代，游戏进程从29步增加到200步。
PPO算法的局限性是只适合于离散策略，因为需要计算同样的动作序列在活动模型与基准模型之间的KL散度。由于自然语言生成任务满足这一约束，在同一序列的多次更新中，反向传播的信息量大幅提升，并且梯度的稳定性比单步更新可靠。

在训练完SFT模型之后，对微调数据集中的每个问题$x$, 利用SFT模型生成两个回答$y_1,y_2$,然后通过人类选择哪个回答更好，按回答好坏分别标记为$y_w\succ y_l\mid x$. 假设有一个最优奖励函数 $r^{\ast }$（完全与人类倾向对齐），则可以将人类倾向的条件概率$p^{\ast }$表示为
$$\begin{array}{c}{p}^{\ast }(y_1\succ y_2\mid x)=\frac{\exp (r^{\ast }(x,y_1))}{\exp (r^{\ast }(x,y_1))+\exp (r^{\ast }(x,y_2))}\end{array}$$
现在假设奖励模型为$r_{\varphi }$，定义损失函数
$$\begin{array}{c}{\mathcal{L}}_R(r_{\varphi },\mathcal{D})=-{\mathbb{E}}_{x\text{~}\mathcal{D},y_w,y_l\text{~}{\pi }_{SFT}(y\mid x)}[\log p_{\varphi }(y_w\succ y_l\mid x)]\end{array}$$
其中
$$p_{\varphi }(y_w\succ y_l\mid x)=\sigma (r_{\varphi }(x,y_w)-r_{\varphi }(x,y_l))$$
注意到$$\sigma (r_w-r_l)=\frac{1}{1+\exp (r_l-r_w)}=\frac{\exp (r_w)}{\exp (r_w)+\exp (r_l)}$$
即
$$\begin{array}{c}{p}_{\varphi }(y_w\succ y_l\mid x)=\frac{\exp (r_{\varphi }(x,y_w))}{\exp (r_{\varphi }(x,y_w))+\exp (r_{\varphi }(x,y_l))}\end{array}$$
与$(1)$的定义一致。只是这里用$r_{\varphi }$代替$r^{\ast }$。
现在回到游戏中，根据论文的定义，我们的目标是最大化以下函数，
$$\begin{array}{c}{L}_t^{CLIP+VF+S}(\theta )={\hat{\mathbb{E}}}_t\left[L_t^{CLIP}(\theta )-c_1{L}_t^{VF}(\theta )+c_{2}S[{\pi }_{\theta }](s_t)\right]\end{array}$$
其中
$$\begin{array}{c}{L}_t^{CLIP}(\theta )={\hat{\mathbb{E}}}_t\left[\min \left(\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{old}}(a_t|s_t)}{\hat{A}}_t,\text{clip}\left(\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{old}}(a_t|s_t)},1-ϵ,1+ϵ\right){\hat{A}}_t\right)\right]\end{array}$$
这里$\pi$是策略函数，PPO更新模型参数是分阶段的，每个阶段更新之前，先保存一个当前模型的副本，记为${\theta }_{old}$，称为参考模型，更新过程中的模型，记为$\theta$，而${\pi }_{\theta }(a_t|s_t)$和${\pi }_{{\theta }_{old}}(a_t|s_t)$就分别代表了当前模型和参考模型对于给定状态$s_t$，给出动作$a_t$的概率。${\hat{A}}_t$表示执行动作$a_t$后的优势函数，所以$\min$中的前一项可以这么理解，假如${\hat{A}}_t>0$，我们希望更新过程中的模型能比参考模型对于动作$a_t$给出更大的概率，而对于${\hat{A}}_t<0$，则希望更新过程中的模型给出更小的概率，总的来说，$L_t^{CLIP}$反映了新模型相对于参考模型的收益。$\min$的第二项，即$\text{clip}$函数，只是为了限制这个比例，防止增益过大造成不良影响。
$A_t$又是怎么来的呢，根据论文的定义，
$$\begin{array}{r}{\hat{A}}_t={\delta }_t+(\gamma \lambda ){\delta }_{t+1}+\cdots +(\gamma \lambda {)}^{T-t+1}{\delta }_{T-1}\\ {\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)\end{array}$$
其中$r_t$来自于游戏中对$s_t$执行动作$a_t$之后得到的奖励，这个奖励可以根据通关时间、新增金币数等信息综合计算，是自定义的激励。$V_t$是模型对当前局面的评估分数。
注意到在LLM里面，$V_t$通常只能由模型估计，只有整个序列结束以后，才有一个人工的奖励，但游戏不一样，每一帧都可以计算奖励，因此$V_t$实际上是有监督信号，这里我们就把$V_t^{\ast }$的目标值设置为当前局面的分数，即$V_t^{\ast }={\Sigma }_{t_i=0}^t{r}_{t_i}$。
$$\begin{array}{r}{L}_t^{VF}(\theta )=(V_t-V_t^{\ast }{)}^2\\ S[{\pi }_{\theta }](s_t)=-{\pi }_{\theta }(s_t)\log {\pi }_{\theta }(s_t)\end{array}$$
分别是行为激励函数的误差和模型基于$s_t$给出的策略的熵。论文中设置$c_2=0.01$，增大策略熵，相当于希望模型给出不确定性更大的策略，即$a_t$的概率分布更均匀。注意这里如果设置$c_2\le 0$，会导致模型很快的陷入局部最优，在游戏中的表现就是一动也不动或按住一个键不放。

第二节 代码实现

完整代码请参考https://github.com/eastonhou/super-mario-ppo.

选择游戏和关卡

在trainer.py文件的最后几行中，可以看到如下代码

ppo = PPO(games.create_mario_profile, dict(world=1, stage=1), 8, 4)

上述示例代码创建了超级玛丽游戏训练对象，其中第一个参数是游戏创建函数，第二个参数设置大关和小关，第三个参数代表每相邻8帧组成一个模型的输入（模型无法从一张静态图推理最佳的执行动作或者计算局面分数，因此需要连续的几帧），具体表现为通道层的层数。第四个参数代表每次跳过4帧，因为AI不需要每帧都进行推理和操作。如果只想看AI如何玩游戏，可以在上述设定游戏参数上，执行

python trainer.py --device cuda

随后可在以下文件夹中找到游戏视频checkpoints文件夹中找到游戏视频。

supermario-game-video

视频上方的橙色格子表示AI对当前画面的预测的各种操作的概率，最终采取的操作通过多项式采样。如果发现游戏停滞不前并且操作锁死，可以通过降低分值系数或增加$c_2$系数（见超参数设置），如果发现操作过于随机，则反方向调整参数。

模型设计

AI模型是采用PyTorch实现的一个简单的卷积网络，由4个卷积层和一个全链接层组成，actor_linear输出策略对数概率，critic_linear输出局势评估分数值。

class GameModel(nn.Module):
    def __init__(self, num_inputs, num_actions) -> None:
        super().__init__()
        # 定义网络结构 
        self.layers = nn.Sequential(
            nn.Conv2d(num_inputs, 32, 3, stride=2, padding=1), nn.Mish(inplace=True),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.Mish(inplace=True),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.Mish(inplace=True),
            nn.Conv2d(128, 32, 3, stride=2, padding=1), nn.Mish(inplace=True),
            nn.Flatten(),
            nn.Linear(1152, 512), nn.LayerNorm(512))
        self.critic_linear = nn.Sequential(nn.Linear(512, 1))
        self.actor_linear = nn.Sequential(
            nn.Linear(512, num_actions),
            nn.LogSoftmax(-1)
        )
	# 前向传播
    def forward(self, input):
        hidden = self.layers(input.float().div(255))
        logits, critic = self.actor_linear(hidden), self.critic_linear(hidden)
        return logits, critic.view(-1)

激励函数的设计

在games.py文件中有Mario和Breakout两个游戏的激励函数定义，以MarioReward为例，定义局面分数为
$$Score=N_{coins}\times 100+X_{mov}-T_{elapsed}+{\delta }_{big}\times 200+{\delta }_{pass}\times 1000-{\delta }_{fail}\times 200$$
通俗地解释，每获得一枚金币得100分，向前走一像素加1分，每花一秒扣1分（为了防止AI躺平），吃了变大蘑菇得200分，通关得1000分，游戏失败扣200分。

class MarioReward(gym.Wrapper):
	...
    def _compute(self, reward, done, info):
        score = info['coins'] * 100
        score += info['x_pos']
        score += (info['time'] - 400)
        if info['status'] == 'big': score += 200
        if done:
            if info['flag_get']: score += 1000
            else: score -= 200
            info['state'] = 'done'
        else:
            info['state'] = 'playing'
        return score

策略梯度的计算

我们定义Loss函数为$-L_t^{CLIP+VF+S}(\theta )$，即(4)式的相反数。因为(4)式是PPO策略的优化目标（越大越好），而pytorch以最小化损失函数为优化目标，因此需要取负号。下面我们分析PPO损失函数的核心代码，可以在trainer.py文件中找到如下代码

    def forward(self, pack, logits, values):
        # 一共有n步
        n = pack['ref_log_probs'].shape[0]
        # 计算每个操作的对数概率
        new_log_probs = logits.log_softmax(-1)
        # 计算当前模型相对参考模型对于每个动作的概率比
        ratio = (new_log_probs[torch.arange(n), pack['actions']] - pack['ref_log_probs']).exp()
        # 剪切概率比，当激励为正时，概率比不大于$1+\epsilon$，当激励函数为负时，概率比不小于$1-\epsilon$
        lb = torch.full_like(ratio, fill_value=0)
        ub = torch.full_like(ratio, fill_value=10000)
        lb[pack['advantages'] < 0] = 1 - self.epsilon
        ub[pack['advantages'] > 0] = 1 + self.epsilon
        # 计算策略收益，对应公式(5)
        actor_loss = -torch.min(
            ratio.mul(pack['advantages']),
            torch.clamp(ratio, lb, ub).mul(pack['advantages'])).mean()
        # 计算价函数函数误差，对应公式(8)
        critic_loss = nn.functional.l1_loss(values, pack['V'], reduction='mean')
        # 计算策略概率自身的熵，对应公式(9)
        entropy_loss = -new_log_probs.exp().mul(new_log_probs).sum(-1).mean()
        loss = actor_loss + critic_loss - self.beta * entropy_loss
        return {
            'loss': loss,
            'actor': actor_loss.item(),
            'critic': critic_loss.item(),
            'entropy': entropy_loss.item(),
            'steps': n
        }

其中误差项entropy_loss的意义在定义公式(9)的地方解释过，读者可以回过去看。
误差项critic_loss，按公式(8)的定义应该是平方误差，这里改为L1误差是因为L1收敛得更平稳一些。
为何需要剪切概率比？ratio是新旧两个模型对于策略$a_t$的预测概率的比值
$$\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{old}}(a_t|s_t)}$$
这个比值可能会很大或很小，如果没加任何约束，训练过程容易导致梯度过大而破坏参数稳定性，这里引用论文中的CLIP逻辑的示意图，表示代码中lb和ub的计算过程。

这里$A$是优势函数，即本文中的$A_t$。$A_t$的计算方式如下

    def precomute(self, ref_logits, ref_values, actions, rewards, scores):
        n = ref_values.shape[0] - 1
        # 计算即时奖励$\delta_t$，对应公式(7)
        delta = rewards + self.gamma * ref_values[1:] - ref_values[:-1]
        # 每个即时奖励乘上时间系数，对应公式(6)
        advantages = delta[None, :].mul(self.coef[:n, :n]).sum(-1)
        ref_log_prob = ref_logits.log_softmax(-1)[torch.arange(n), actions]
        # 获得$V_t$的目标值，这里以局面分数作为局面评估函数的目标
        V = scores.float()
        return V, advantages, ref_log_prob

ref_log_prob是参考模型的策略概率${\pi }_{{\theta }_{old}}(a_t|s_t)$，scores相当于即时奖励rewards的累加，由游戏奖励函数定义。这两个返回值与$A_t$无关，但是在计算损失函数时会用到，所以这里顺便计算一下。

超参数设置

	论文中的参数值	游戏中的参考值	参数说明
$\gamma$	0.99	1	价值分数折扣率
$\lambda$	0.95	0.975	激励衰减率
$ϵ$	0.2	0.2	优势函数剪切窗口系数
$c_1$	1 (MSE)	1(L1)	价值损失函数系数
$c_2$	0.01	0.1$\alpha$	探索系数
$l_r$	$2.5\times 10^{-4}\alpha$ (Adam)	$2.5\times 10^{-4}\alpha$ (AdamW)	学习率

其中$\alpha$是随训练进程从1向0逐渐减小的控制参数，游戏中设置每次迭代衰减率为0.01。其余参数可以在Loss类的的定义中找到相应设置。

class Loss(nn.Module):
    def __init__(self, gamma=1, lmbda=0.975, epsilon=0.2, c2=0.1) -> None:
        super().__init__()
        self.gamma = gamma
        self.lmbda = lmbda
        self.c2 = c2
        self.epsilon = epsilon

注意$c_2$和学习率是随时间衰减的，50个迭代后下降至60%。

训练日志

机器配置：13900K+RTX3090
训练时间：约20分钟

训练曲线的解释
actor: 优势函数的相反数，越小越好，反映AI操作的熟练程度。
critic: 价值函数的精度，越小越好，虽然曲线看起来不收敛，但价值函数的具体值与游戏进程有关，并且策略梯度的计算主要依赖邻近状态的差值，与绝对数值关系不大。
entropy: AI对操作的困惑度，一般在0.2~0.6之间比较好，太小会导致过早收敛，探索不到更好的全局目标，太大则操作不稳定，容易挂。
score: 游戏中所得分数，越高越好。
steps: 游戏结束时间，在能通关的前提下，越小越好。

总结

本文主要是为了复原PPO的算法实现，但PPO算法本身未必是最适合游戏的增强学习策略，读者可以自行修改策略函数，超参数，以及奖励函数，构造出更智能的游戏模型。PPO的超参数调整过程很玄学，各参数之间必须保持平衡，否则容易不收敛，调参时建议逐个参数测试，做好实验笔记。有什么想法欢迎交流。
邮箱：easton.hou@outlook.com。