Added loss vizualisation

reinforce_cartpole.py

@@ -28,9 +28,11 @@ optimizer = optim.Adam(model.parameters(), lr=5e-3)
 
 # Keep track of the number of rewards for each episodes
 rewardsByEpisode = []
+lossesByEpisode = []
 
 # training loop
 for episode in range(500):
+    print("Episode's number {0} / {1}".format(episode, 500))
     # reset the environment
     state = env.reset()
     log_probs = []
@@ -69,18 +71,31 @@ for episode in range(500):
         model_loss.append(-log_prob * return_)
     model_loss = torch.cat(model_loss).sum()
 
+    lossesByEpisode.append(model_loss.item())
+
+    
     # update the model
     optimizer.zero_grad()
     model_loss.backward()
     optimizer.step()
-    print("N° de l'épisode :", episode)
-    print("Nombre de rewards :", len(rewards))
+
 
 # X axis :
 x = list(range(len(rewardsByEpisode)))
 
+# Plot Rewards by Episodes
+plt.figure("Figure 1")
 plt.xlabel('Episodes N°')
 plt.ylabel('Number of rewards given')
-plt.plot(x, rewardsByEpisode, '--')
+plt.plot(x, rewardsByEpisode)
+plt.savefig('reinforce_cartpole_reward.png')
+plt.show()
+
+# Plot losses by Episodes
+plt.figure("Figure 2")
+plt.xlabel("Episodes N°")
+plt.ylabel("Loss")
+plt.plot(x, lossesByEpisode)
+plt.title("Loss by episodes")
+plt.savefig("reinforce_cartpole_loss.png")
 plt.show()
-plt.savefig('reinforce_cartpole.png')