diff --git a/entities/player.py b/entities/player.py
index 0a38e3c..904ff5f 100644
--- a/entities/player.py
+++ b/entities/player.py
@@ -161,7 +161,7 @@ class Player(pygame.sprite.Sprite):
             self.old_state_features = self.state_features
 
             self.reward = self.stats.exp
-                # + self.stats.health/self.stats.stats['health'] - 1\
+                + self.stats.health/self.stats.stats['health'] - 1
                 # - nearest_dist/np.sqrt(np.sum(self.map_edge))
 
         self.state_features = [
diff --git a/level.py b/level.py
index c61a4db..cfb0112 100644
--- a/level.py
+++ b/level.py
@@ -294,7 +294,6 @@ class Level:
             if player.is_dead():
                 print(f"\nPlayer {player.player_id} is dead\n")
                 player.stats.exp = -10
-                player.update()
                 self.dead_players[player.player_id] = player.is_dead()
 
         self.done = True if (self.dead_players.all() == 1
diff --git a/main.py b/main.py
index 83982e0..e474a41 100644
--- a/main.py
+++ b/main.py
@@ -26,6 +26,7 @@ def main():
 
     chkpt_path, figure_path = folder_struct.setup_dirs()
         
+    # Setup AI params
     n_episodes = parsed_args.n_episodes
     episode_length = parsed_args.ep_length
     n_agents = parsed_args.n_agents
@@ -38,8 +39,6 @@ def main():
 
     show_pygame = parsed_args.show_pg
 
-    # Setup AI metrics
-
     # Setup parameter monitoring
     score_history = np.zeros(
         shape=(parsed_args.n_agents, parsed_args.n_episodes))
@@ -63,9 +62,6 @@ def main():
 
     time_alive = np.zeros(shape=(parsed_args.n_agents,
                                  parsed_args.n_episodes))
-    # score_history, best_score, actor_loss, critic_loss, total_loss, entropy, advantage  = metrics.generate(parsed_args)
-    
-    
     game = Pneuma(show_pg=show_pygame, n_players=parsed_args.n_agents)
 
     print("Initializing agents ...")
@@ -202,9 +198,10 @@ def main():
         metrics.plot_parameter('advantage', advantage, parsed_args.n_agents, figure_path)
 
         metrics.plot_avg_time(time_alive, parsed_args.n_agents, figure_path)    
+
     # End of training session
     print("End of episodes.\
-        \nExiting game...")
+        \n Saving models and exiting game...")
 
     # Save models
     player.agent.save_models(