情報系院生のノート

情報技術系全般,自分用メモを公開してます。

強化学習の報酬のグラフを良い感じに書く

完成図

f:id:tontainoti:20210331174324p:plain
報酬のグラフ

こんな感じで2つのアルゴリズムを比較できるように作りました。もちろん、1つでも使えます。

(報酬が離散的すぎてやや例としては悪いですが…)

想定するデータ

なんでも良いのですが、今回はtensorboardからDownloadした.csvを前提に作業を進めます。

以下のような感じでロギングしたデータを使います。

writer.add_scalar("data/reward",reward, step)

コード

説明

グラフのスムージングには指数移動平均(EMA)を用います。(厳密にはpandasのewm(Exponentially weighted moving average)は指数移動平均ではないみたいですが、おおよその報酬の動きが分かれば良いのでまあいいでしょう)

すべてmatplotlibで完結させても良いのですが、seabornが個人的に好きなので、sns.lineplot使ってます。 seaborn入れたくない人は、matplotlibのplot関数で色を指定すると同じようなグラフが出来上がると思います。

定数

SPAN = 50
ALPHA = 0.25

# PDDDQNのcsv
df_dqn = pd.read_csv(path-to-csv)
df_dqn['Algorithm'] = np.tile(['PDDDQN'], len(df_dqn))
df_dqn['Reward'] = df_dqn['Value'].ewm(span=SPAN).mean()

# A2Cのcsv
df_a2c = pd.read_csv(path-to-csv)
df_a2c['Algorithm'] = np.tile(['A2C'], len(df_a2c))
df_a2c['Reward'] = df_a2c['Value'].ewm(span=SPAN).mean()


df = pd.concat([df_dqn, df_a2c])
df.reset_index()

# 後ろの薄いグラフ
plt.plot(df_dqn['Step'], df_dqn['Value'], alpha=ALPHA)
plt.plot(df_a2c['Step'], df_a2c['Value'], alpha=ALPHA)


# 前の指数平均線
ax = sns.lineplot(data=df, x='Step', y='Reward', hue='Algorithm')

# 高さ調整
ax.set_ylim([0.0, 0.35])

# legendの位置
ax.legend(loc=1)

ax.set_title('Experiment 1')

plt.savefig('rewards.png', dpi=500)