本文使用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(θ)−c1LtVF(θ)+c2S[πθ](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=0trti。
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文件的最后几行中,可以看到如下代码
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输出局势评估分数值。
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')
# 计算策略概率自身的熵,对应公式(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的计算方式如下
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无关,但是在计算损失函数时会用到,所以这里顺便计算一下。
超参数设置
论文中的参数值 | 游戏中的参考值 | 参数说明 | |
---|---|---|---|
γ \gamma γ | 0.99 | 1 | 价值分数折扣率 |
λ \lambda λ | 0.95 | 0.975 | 激励衰减率 |
ϵ \epsilon ϵ | 0.2 | 0.2 | 优势函数剪切窗口系数 |
c 1 c_1 c1 | 1 (MSE) | 1(L1) | 价值损失函数系数 |
c 2 c_2 c2 | 0.01 | 0.1 α \alpha α | 探索系数 |
l r l_r lr | 2.5 × 1 0 − 4 α 2.5\times10^{-4}\alpha 2.5×10−4α (Adam) | 2.5 × 1 0 − 4 α 2.5\times10^{-4}\alpha 2.5×10−4α (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 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。