本文使用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

x

x, 利用SFT模型生成两个回答

y

1

,

y

2

y_1,y_2

y1​,y2​,然后通过人类选择哪个回答更好，按回答好坏分别标记为

y

w

≻

y

l

∣

x

y_w \succ y_l \mid x

yw​≻yl​∣x. 假设有一个最优奖励函数

r

∗

r^*

r∗（完全与人类倾向对齐），则可以将人类倾向的条件概率

p

∗

p^*

p∗表示为

p

∗

(

y

1

≻

y

2

∣

x

)

=

exp

⁡

(

r

∗

(

x

,

y

1

)

)

exp

⁡

(

r

∗

(

x

,

y

1

)

)

+

exp

⁡

(

r

∗

(

x

,

y

2

)

)

\begin{equation} p^*(y_1\succ y_2 \mid x)=\frac{\exp(r^*(x,y_1))}{\exp(r^*(x,y_1))+\exp(r^*(x,y_2))} \end{equation}

p∗(y1​≻y2​∣x)=exp(r∗(x,y1​))+exp(r∗(x,y2​))exp(r∗(x,y1​))​​​

现在假设奖励模型为

r

ϕ

r_\phi

rϕ​，定义损失函数

L

R

(

r

ϕ

,

D

)

=

−

E

x

~

D

,

y

w

,

y

l

~

π

S

F

T

(

y

∣

x

)

[

log

⁡

p

ϕ

(

y

w

≻

y

l

∣

x

)

]

\begin{equation} \mathcal{L}_R(r_\phi,\mathcal{D})=-\mathbb{E}_{x\text{\textasciitilde}\mathcal{D},y_w,y_l\text{\textasciitilde}\pi_{SFT}(y\mid x)}[\log{p_\phi(y_w\succ y_l \mid x)}] \end{equation}

LR​(rϕ​,D)=−Ex~D,yw​,yl​~πSFT​(y∣x)​[logpϕ​(yw​≻yl​∣x)]​​

其中

p

ϕ

(

y

w

≻

y

l

∣

x

)

=

σ

(

r

ϕ

(

x

,

y

w

)

−

r

ϕ

(

x

,

y

l

)

)

p_\phi(y_w\succ y_l \mid x)=\sigma(r_\phi(x,y_w)-r_\phi(x,y_l))

pϕ​(yw​≻yl​∣x)=σ(rϕ​(x,yw​)−rϕ​(x,yl​))

注意到

σ

(

r

w

−

r

l

)

=

1

1

+

exp

⁡

(

r

l

−

r

w

)

=

exp

⁡

(

r

w

)

exp

⁡

(

r

w

)

+

exp

⁡

(

r

l

)

\sigma(r_w-r_l)=\frac{1}{1+\exp(r_l-r_w)}=\frac{\exp(r_w)}{\exp(r_w)+\exp(r_l)}

σ(rw​−rl​)=1+exp(rl​−rw​)1​=exp(rw​)+exp(rl​)exp(rw​)​

即

p

ϕ

(

y

w

≻

y

l

∣

x

)

=

exp

⁡

(

r

ϕ

(

x

,

y

w

)

)

exp

⁡

(

r

ϕ

(

x

,

y

w

)

)

+

exp

⁡

(

r

ϕ

(

x

,

y

l

)

)

\begin{equation} p_\phi(y_w\succ y_l \mid x)=\frac{\exp(r_\phi(x,y_w))}{\exp(r_\phi(x,y_w))+\exp(r_\phi(x,y_l))} \end{equation}

pϕ​(yw​≻yl​∣x)=exp(rϕ​(x,yw​))+exp(rϕ​(x,yl​))exp(rϕ​(x,yw​))​​​

与

(

1

)

(1)

(1)的定义一致。只是这里用

r

ϕ

r_\phi

rϕ​代替

r

∗

r^*

r∗。

现在回到游戏中，根据论文的定义，我们的目标是最大化以下函数，

L

t

C

L

I

P

+

V

F

+

S

(

θ

)

=

E

^

t

[

L

t

C

L

I

P

(

θ

)

−

c

1

L

t

V

F

(

θ

)

+

c

2

S

[

π

θ

]

(

s

t

)

]

\begin{equation} 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{equation}

LtCLIP+VF+S​(θ)=E^t​[LtCLIP​(θ)−c1​LtVF​(θ)+c2​S[πθ​](st​)]​​

其中

L

t

C

L

I

P

(

θ

)

=

E

^

t

[

min

⁡

(

π

θ

(

a

t

∣

s

t

)

π

θ

o

l

d

(

a

t

∣

s

t

)

A

^

t

,

clip

(

π

θ

(

a

t

∣

s

t

)

π

θ

o

l

d

(

a

t

∣

s

t

)

,

1

−

ϵ

,

1

+

ϵ

)

A

^

t

)

]

\begin{equation} 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-\epsilon,1+\epsilon\right)\hat{A}_t \right)} \right] \end{equation}

LtCLIP​(θ)=E^t​[min(πθold​​(at​∣st​)πθ​(at​∣st​)​A^t​,clip(πθold​​(at​∣st​)πθ​(at​∣st​)​,1−ϵ,1+ϵ)A^t​)]​​

这里

π

\pi

π是策略函数，PPO更新模型参数是分阶段的，每个阶段更新之前，先保存一个当前模型的副本，记为

θ

o

l

d

\theta_{old}

θold​，称为参考模型，更新过程中的模型，记为

θ

\theta

θ，而

π

θ

(

a

t

∣

s

t

)

\pi_\theta(a_t|s_t)

πθ​(at​∣st​)和

π

θ

o

l

d

(

a

t

∣

s

t

)

\pi_{\theta_{old}}(a_t|s_t)

πθold​​(at​∣st​)就分别代表了当前模型和参考模型对于给定状态

s

t

s_t

st​，给出动作

a

t

a_t

at​的概率。

A

^

t

\hat{A}_t

A^t​表示执行动作

a

t

a_t

at​后的优势函数，所以

min

⁡

\min

min中的前一项可以这么理解，假如

A

^

t

>

0

\hat{A}_t>0

A^t​>0，我们希望更新过程中的模型能比参考模型对于动作

a

t

a_t

at​给出更大的概率，而对于

A

^

t

<

0

\hat{A}_t<0

A^t​<0，则希望更新过程中的模型给出更小的概率，总的来说，

L

t

C

L

I

P

L_t^{CLIP}

LtCLIP​反映了新模型相对于参考模型的收益。

min

⁡

\min

min的第二项，即

clip

\text{clip}

clip函数，只是为了限制这个比例，防止增益过大造成不良影响。

A

t

A_t

At​又是怎么来的呢，根据论文的定义，

A

^

t

=

δ

t

+

(

γ

λ

)

δ

t

+

1

+

⋯

+

(

γ

λ

)

T

−

t

+

1

δ

T

−

1

δ

t

=

r

t

+

γ

V

(

s

t

+

1

)

−

V

(

s

t

)

\begin{align} \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{align}

A^t​=δt​+(γλ)δt+1​+⋯+(γλ)T−t+1δT−1​δt​=rt​+γV(st+1​)−V(st​)​​

其中

r

t

r_t

rt​来自于游戏中对

s

t

s_t

st​执行动作

a

t

a_t

at​之后得到的奖励，这个奖励可以根据通关时间、新增金币数等信息综合计算，是自定义的激励。

V

t

V_t

Vt​是模型对当前局面的评估分数。

注意到在LLM里面，

V

t

V_t

Vt​通常只能由模型估计，只有整个序列结束以后，才有一个人工的奖励，但游戏不一样，每一帧都可以计算奖励，因此

V

t

V_t

Vt​实际上是有监督信号，这里我们就把

V

t

∗

V_t^*

Vt∗​的目标值设置为当前局面的分数，即

V

t

∗

=

Σ

t

i

=

0

t

r

t

i

V_t^*=\Sigma_{t_i=0}^{t}{r_{t_i}}

Vt∗​=Σti​=0t​rti​​。

L

t

V

F

(

θ

)

=

(

V

t

−

V

t

∗

)

2

S

[

π

θ

]

(

s

t

)

=

−

π

θ

(

s

t

)

log

⁡

π

θ

(

s

t

)

\begin{align} L_t^{VF}(\theta)=(V_t-V_t^*)^2\\ S[\pi_\theta](s_t)=-\pi_\theta(s_t)\log{\pi_\theta (s_t)} \end{align}

LtVF​(θ)=(Vt​−Vt∗​)2S[πθ​](st​)=−πθ​(st​)logπθ​(st​)​​

分别是行为激励函数的误差和模型基于

s

t

s_t

st​给出的策略的熵。论文中设置

c

2

=

0.01

c_2=0.01

c2​=0.01，增大策略熵，相当于希望模型给出不确定性更大的策略，即

a

t

a_t

at​的概率分布更均匀。注意这里如果设置

c

2

≤

0

c_2\le0

c2​≤0，会导致模型很快的陷入局部最优，在游戏中的表现就是一动也不动或按住一个键不放。

第二节 代码实现

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

选择游戏和关卡

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

<code>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

c_2

c2​系数（见超参数设置），如果发现操作过于随机，则反方向调整参数。

模型设计

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

<code>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为例，定义局面分数为

S

c

o

r

e

=

N

c

o

i

n

s

×

100

+

X

m

o

v

−

T

e

l

a

p

s

e

d

+

δ

b

i

g

×

200

+

δ

p

a

s

s

×

1000

−

δ

f

a

i

l

×

200

Score=N_{coins}\times100+X_{mov}-T_{elapsed}+\delta_{big}\times200+\delta_{pass}\times1000-\delta_{fail}\times200

Score=Ncoins​×100+Xmov​−Telapsed​+δbig​×200+δpass​×1000−δfail​×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

C

L

I

P

+

V

F

+

S

(

θ

)

-L_t^{CLIP+VF+S}(\theta)

−LtCLIP+VF+S​(θ)，即(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')code>

# 计算策略概率自身的熵，对应公式(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

a_t

at​的预测概率的比值

π

θ

(

a

t

∣

s

t

)

π

θ

o

l

d

(

a

t

∣

s

t

)

\frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}

πθold​​(at​∣st​)πθ​(at​∣st​)​

这个比值可能会很大或很小，如果没加任何约束，训练过程容易导致梯度过大而破坏参数稳定性，这里引用论文中的CLIP逻辑的示意图，表示代码中lb和ub的计算过程。

这里

A

A

A是优势函数，即本文中的

A

t

A_t

At​。

A

t

A_t

At​的计算方式如下

<code> 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是参考模型的策略概率

π

θ

o

l

d

(

a

t

∣

s

t

)

\pi_{\theta_{old}}(a_t|s_t)

πθold​​(at​∣st​)，scores相当于即时奖励rewards的累加，由游戏奖励函数定义。这两个返回值与

A

t

A_t

At​无关，但是在计算损失函数时会用到，所以这里顺便计算一下。

超参数设置

其中

α

\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

c_2

c2​和学习率是随时间衰减的，50个迭代后下降至60%。

训练日志

机器配置：13900K+RTX3090

训练时间：约20分钟

训练曲线的解释

actor: 优势函数的相反数，越小越好，反映AI操作的熟练程度。

critic: 价值函数的精度，越小越好，虽然曲线看起来不收敛，但价值函数的具体值与游戏进程有关，并且策略梯度的计算主要依赖邻近状态的差值，与绝对数值关系不大。

entropy: AI对操作的困惑度，一般在0.2~0.6之间比较好，太小会导致过早收敛，探索不到更好的全局目标，太大则操作不稳定，容易挂。

score: 游戏中所得分数，越高越好。

steps: 游戏结束时间，在能通关的前提下，越小越好。

总结

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

邮箱：easton.hou@outlook.com。