Highly customized seaborn figure

beautiful_seaborn_fig.py

from toolbox.process_data import parse

data3 = []

a2c_data3 = parse("/home/zhpeng/OLD-novel-rl-archive-20200210/0206-a2c-baseline-large-Walker2d-v3-Walker2d-v3-50000000ts.pkl")
a2c_data3["label"] = "A2C"
a2c_data3.episode_reward_mean *= 5
data3.append(a2c_data3)

ppo_data3 = parse("/home/zhpeng/0119-ppo-baseline-walker-Walker2d-v3-50000000ts.pkl")
ppo_data3["label"] = "PPO"
# ppo_data3 = ppo_data3[ppo_data3.timesteps_total < 5e6]
ppo_data3.episode_reward_mean *= 5
data3.append(ppo_data3)

d = parse("/home/zhpeng/ray_results/0122-dece-normalize_advantage-normalize_advantage")
d["label"] = "DiCE"
data3.append(d)

d = parse("/home/zhpeng/ray_results/dice_0419ablation-Walker2d-v3")
d["label"] = "WRONG"
d.loc[(d.only_tnb==False) & (d.use_bisector==False), "label"] = "w/o DR"
d.loc[(d.only_tnb==True) & (d.use_bisector==True), "label"] = "w/o CE"
data3.append(d)

data3 = pd.concat(data3)

data3 = data3[data3.label != "WRONG"]

data3.episode_reward_mean /= 5

# data3 = data3.sort_values("label")

mask = (data3["label"]=="PPO") | (data3["label"]=="A2C")
data3["style"] = ""

# The order make sense
data3.loc[~mask, "style"] = 1
data3.loc[mask, "style"] = 2

plt.figure(figsize=(6, 3))
sns.set("notebook", "white")

ax = sns.lineplot(
    x="timesteps_total",
    y="episode_reward_mean", 
    style='style',
    hue='label',
    data=data3,
    err_kws={'alpha': 0.15},
    ci='sd'
)

ax.ticklabel_format(style='sci', scilimits=(0,0), axis='x')
ax.set_xlabel("Sampled Steps")
ax.set_ylabel("Episode Reward")
ax.set_title("Walker2d-v3")
plt.legend(labels=[
    "A2C",
    "PPO",
    "DiCE",
    "w/o CE",
    "w/o DR"
], prop={'size': "x-small"})

plt.savefig('DICE-ICML20-fig/ablate-walker-formal.pdf', format='pdf', dpi=300, bbox_inches = "tight")

This script results a publication-level figure: