正梯度仍然使得policy的probability增大

现象

使用policy gradient算法写了一个小模型,训练过程中发现模型很快会收敛到一个非常糟糕的结果,理论上不应该呀,因为按照policy gradient算法原理来说,对于某个action,如果你给了负的reward,那么模型会减小该action出现的概率,但我的实验表明,模型不仅没减小该action的概率,甚至还会增加它的概率,最终它的概率甚至会等于1。理论上,即使这个不好的action概率为1了,但是reward是负的,每次迭代应该还是会降低其概率,但我这个实验发现并没有,反而会继续降低另外一个好的action的概率——即使它已经无限逼近0了。

排查

首先我想,既然坏的action概率在增加,有没有可能是因为reward是正数(本来应该给负数),但是经过简单的排查后发现,该action的reward就是负数。

然后考虑到,policy gradient 对采样的样本action会无脑增大其发生概率(前提是该action的reward是正的),然而采样所使用的概率也是这个概率,那么就可能形成一个恶行循环:梯度上升导致概率增大->概率增大导致其更容易被采样->采样后概率继续增大…。但经过计算发现,即使它的概率会恶性增大,但如果采样的足够多,总是会采样到其他reward更大的action上,那么它的概率还是会减小很多。另一点就是经过排查,发现这个action的reward确实是负的,应该不会出现这个情况。

第三个想法是,因为这是一个神经网络,本质上就是一个函数,假如这个函数有图像的话,现将其简化想象成一个二维的曲线,那么上面的问题就类比于:如果我采样到上面某个点,经过计算发现下一次迭代后,这个点的位置应该下降(本质上是函数发生了变化,原本的函数更新为另一个新的函数),但是实验却发现下一次迭代后它还上升了。一个很直接的想法就是,有没有可能,如果单独更新这一个点时,它确实会下降,但实际上,我同时更新了多个点,在更新其他点时,更新的参数影响了这个点的值,那么它就有可能不降反升。理论上,这也确实是非常有可能的,但它发生的概率应该要远小于50%,否则模型只会越训练越差。但上述实验发现,好像这个值又接近50%,这就让人匪夷所思了。而且,理论上,就算出现这种问题,在迭代次数足够多后,这种情况应该越来越少。所以这个假设也被排除。

实验到此已经很晚了,我仍然没想通什么原因。于是我想会不会就算训练的时间不够导致的,于是我同时运行好几个相同的模型,其中只有部分有修改,例如有些使用 discount reward,有些使用 total reward,有些则更换优化器等,跑一个晚上看看结果。

如何使用梯度判断模型中间量的更新方向?

在pytorch中,只有叶子节点的变量(学习量)才可以直接获取梯度,但在上述问题中,我发现模型输出的某个概率增大了,我怀疑这里的梯度有问题,那么我想查看这个输出概率的梯度怎么办?很明显直接看是看不了的,因为它只是一个中间量,虽然会计算它的梯度,但并不会在backward后保存下来,因为并不需要真正更新它。

pytorch提供了 register_hook() 函数来获取或修改这些中间量的梯度:

output = model(input)
...
# 如果同一个变量注册了多个 hook,则会顺序执行它们。
output.register_hook(lambda grad: print(grad))  # 读取梯度
output.register_hook(lambda grad: grad * 2)  # 也可以直接修改梯度
...
loss = loss_model(output)
loss.backward()  # hook中的函数会在这里计算完梯度后执行
...

通过这种方法,你就可以知道下一次迭代时是会增加这个输出还是会减少这个输出了。由于迭代后的值等于原值减去梯度,如果梯度是正数,则会减小,反之会增加。下面是一个register_hook()使用方法的例子:

a = torch.tensor(2., requires_grad=True)
b = 5 * a
c = b * b

b.register_hook(lambda g: print(f'b grad: {g}'))  # b grad: 20.0,因为用c对b求偏导为 2b, 而b又等于 5*2,所以b的梯度为 2 * 5 * 2
c.backward()

一个具体的例子:

...
action_probs = actor(torch.tensor(obs, dtype=torch.float))
action = torch.distributions.Categorical(action_probs).sample().item() if np.random.random() < 1 else np.random.randint(0, 2)
...
action_probs.register_hook(lambda g: print(g))
loss = -reward * torch.log(action_probs[action])
loss.backward()

程序运行后查看 action_probs 在 action 位置上的梯度确实是大于0的,由于log函数的导数为 1/x,容易得到等式 -reward * 1/action_probs[action] = grad。由此可以算出 reward = -grad * action_probs[action]。

为什么还要单独去算一遍reward呢?不是程序给定的吗?

因为reward可能会经过discount运算,使得其值具体是多少不确定,甚至不确定正负负号。再者就算检验用了。

原因

多个对比实验结果非常明显,使用adam优化器模型获得的total reward确有上升,而使用SGD的模型就陷入了上面说的问题,模型永远只选择一个策略,且不再更新。问题明显出现在学习率上。

于是,我将SGD的学习率从0.01降低到0.001,发现模型确实又可以训练了。以下为解释:

学习率过大模型在迭代时可能跳过极值点,到达一个比当前位置还大的点,显然,下图中,红色的点表示迭代之前,它的梯度是负的,所以它的x轴应该向正方向移动一点,期望上,它的函数值y应该减小,最终逼近极值点,但由于学习率过大,导致它向x轴正方向移动得太多了,反而使得函数值y比以前还大。然后在下一次迭代后,又跳回到极值点的左边,然后反复横跳,无法收敛。这大概就算我目前遇到的原因了。

解决方法是调小学习率,使得它在梯度较大的位置时只移动很小一步——这正是adam优化器做的事

关于adam优化器见 https://blog.woyou.cool/posts/5363/

另一个问题是,既然使用大的学习率,参数已经可以在极值点来回跳跃了,是不是可以假设函数已经找到一个局部最优解了,那么使用adam或者调小学习率不也是使得参数在这个极值点跳来跳去,只不过幅度没那么大了而已,那函数应该还是收敛在这个局部最优解下,那么模型应该同样练不起来才是。这个想法或许对于非常简单的模型是有可能的,但我这个模型是MLP,已经不简单了。首先,在调整学习率后,它不一定还会收敛到这个极值点,有可能收敛到另一个更好的极值点,再者就算对于MLP来说,其实本质上是很多个函数嵌套的复合函数,内部的函数值改变一点,都有可能使得外部的函数值发生很大变化,所以即便内部函数在同一个极值点处收敛,但它们左右震荡的幅度不同,模型整体的表现可能大不相同。

Leave a Comment