python的EMA曲线平滑方法

写在前面

最近用到了强化学习（DQN），可这东西训练的结果实在是不够稳定，reward波动性极强。肉眼是能看出来reward有上升趋势的，但是不是很明显，还是得做一下曲线平滑。

网上查了很多ema的方法，但是大多都是用在模型当中的，我就是只想做个简单的数据平滑而已啊，没必要大动干戈，如果有什么一行的代码给我最好不过了。其实我也查了，scipy的包中有实现数据平滑的方法，但我试了一下，还挺麻烦的。不如直接写一个简单程序造福大众。

因为我的主程序使用pytorch写的，所以代码中我也用了pytorch，有需要的朋友需要手动变更为numpy的，其实也很简单。

我定义了一个ema函数，其实ema的原理是

ema数据=衰减率*当前数据+(1-衰减率)*下个数据。

知道原理之后我们就可以通过for循环很简单的实现ema了。

import matplotlib.pyplot as plt
import torch


def ema(data, decay=0.85):
    new_data = torch.zeros(data.shape[0])
    new_data[0] = torch.mean(data[:2])
    for idx in range(len(data) - 1):
        new_data[idx+1] = decay * new_data[idx] + (1 - decay) * data[idx + 1]
    return new_data


a = torch.randn(100)
b = ema(a)


plt.figure()
plt.plot(a, label='real')
plt.plot(b, label='ema data')
plt.legend(loc='best')
plt.show()

下面是效果图，设置deacy的大小可以调整平滑的程度。这里我设置的0.85，你可以自己调整一下试试。

再贴一张强化学习的reward图

左边就是DQN的reward数据平滑了，右边是我其他程序预测的效果图，请忽略。