from baselines.common import plot_util as pu
results1 = pu.load_results('../Documents/PhD/allresults/IceHockey_ppo2_4c/ppo2_single-0')
results2 = pu.load_results('../Documents/PhD/allresults/IceHockey_garl_ppo2')
#print(len(results))

import matplotlib.pyplot as plt
import pandas as pd

r1 = results1[0]
r2 = results2[0]
#pu.plot_results_garl(results)
#print(type(r1.progress.eprewmean))
r1s = r1.progress['misc/total_timesteps']
r2s = r2.progress['misc/total_timesteps']
r2s_2 = r2s.append(r1s[r1s > r2s[r2s.size-1]])
r2_2 = pd.Series([r2.progress.eprewmean[r2.progress.eprewmean.size-1] for i in range(int(r2s_2.size-r2.progress.eprewmean.size))])
r2_3 = r2.progress.eprewmean.append(r2_2)
plt.plot(r1.progress['misc/total_timesteps'], r1.progress.eprewmean, label='single-agent')
plt.plot(r2s_2, r2_3, label='group-agent')
plt.xlim([1.5e4, 3.5e5])
plt.title('PPO')
plt.xlabel('timesteps')
plt.ylabel('reward')
plt.legend()
plt.show()
