Updated plot

CNN_PPO/ppo_template_cnn.py

@@ -1,123 +1,146 @@
-
-import gymnasium as gym
-import sys
-import matplotlib.pyplot as plt
-import ale_py
-from ppo_helpers_cnn import *
-from gymnasium.spaces import Box
-import cv2
-
-def preprocess(obs):
-    # Convert to grayscale
-    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
-    # Resize
-    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
-    # Add channel dimension and normalize
-    return np.expand_dims(obs, axis=0).astype(np.float32) / 255.0
-
-def main() -> int:
-    # env = gym.make("ALE/SpaceInvaders-v5", render_mode='human')
-    #env = gym.make("ALE/Pacman-v5", render_mode="human")
-    env = gym.make("ALE/Pacman-v5")
-
-    episode = 0
-    total_return = 0
-    ep_return = 0
-    steps = 100
-    batches = 100
-
-    print("Observation space:", env.observation_space)
-    print("Action space:", env.action_space)
-    """
-    agent = Agent(obs_space=env.observation_space, action_space=env.action_space,
-                  hidden=64, lr=3e-4, gamma=0.99, clip_coef=0.2,
-                  entropy_coef=0.01, value_coef=0.5, seed=70,
-                  batch_size = 64, ppo_epochs = 4, lam = 0.95)
-    
-    """
-    obs, _ = env.reset()
-    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
-
-    agent = Agent(obs_space=dummy_obs_space, action_space=env.action_space,
-                  hidden=64, lr=3e-4, gamma=0.99, clip_coef=0.2,
-                  entropy_coef=0.01, value_coef=0.5, seed=70,
-                  batch_size = 64, ppo_epochs = 4, lam = 0.95)
-
-    # === Return-Based Scaling stats ===
-    r_mean, r_var = 0.0, 1e-8
-    g2_mean = 1.0
-
-    agent.r_var = r_var
-    agent.g2_mean = g2_mean
-
-    try:
-        obs, info = env.reset(seed=42)
-        state = preprocess(obs)
-
-        loss_history = []
-        reward_history = []
-
-        for update in range(1, batches + 1):
-            for t in range(steps):
-                action, logp, value = agent.choose_action(state)
-                next_obs, reward, terminated, truncated, info = env.step(action)
-                done = terminated or truncated
-                next_state = preprocess(next_obs)
-
-                agent.remember(state, action, reward, done, logp, value, next_state)
-
-                ep_return += reward
-                state = next_state
-
-                if done:
-                    episode += 1
-                    total_return += ep_return
-                    print(f"Episode {episode} return: {ep_return:.2f}")
-                    ep_return = 0
-                    obs, info = env.reset()
-                    state = preprocess(obs)
-
-            avg_loss = agent.update_reward_gradient_clipping()
-            loss_history.append(avg_loss)
-
-            avg_ret = (total_return / episode) if episode else 0
-            reward_history.append(avg_ret)
-            print(f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}, avg_loss={avg_loss:.4f}")
-
-        fig = plt.figure()
-
-        ax1 = plt.subplot(321)
-        ax1.plot(agent.sigma_history, label="Return σ")
-        ax1.set_xlabel("PPO Update")
-        ax1.set_ylabel("σ (Return Std)")
-
-        ax2 = plt.subplot(322)
-        ax2.plot(loss_history, label="Avg Loss")
-        ax2.set_ylabel("Average PPO Loss")
-        ax2.set_xlabel("PPO Update")
-
-        ax3 = plt.subplot(323)
-        ax3.plot(reward_history, label="Reward")
-        ax3.set_ylabel("Reward")
-        ax3.set_xlabel("PPO Update")
-
-        fig.suptitle("PPO Training Stability")
-        fig.tight_layout()
-        plt.show()
-
-
-
-
-    except Exception as e:
-        print(f"Error: {e}", file=sys.stderr)
-        return 1
-    finally:
-        avg = total_return / episode if episode else 0
-        print(f"\nEpisodes: {episode}, Avg return: {avg:.3f}")
-        env.close()
-
-    return 0
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+
+import gymnasium as gym
+import sys
+import matplotlib.pyplot as plt
+import ale_py
+from ppo_helpers_cnn import *
+from gymnasium.spaces import Box
+import cv2
+
+def preprocess(obs):
+    # Convert to grayscale
+    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
+    # Resize
+    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
+    # Add channel dimension and normalize
+    return np.expand_dims(obs, axis=0).astype(np.float32) / 255.0
+
+def main() -> int:
+    # env = gym.make("ALE/SpaceInvaders-v5", render_mode='human')
+    #env = gym.make("ALE/Pacman-v5", render_mode="human")
+    env = gym.make("ALE/Pacman-v5")
+
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 100
+    batches = 100
+
+    print("Observation space:", env.observation_space)
+    print("Action space:", env.action_space)
+    """
+    agent = Agent(obs_space=env.observation_space, action_space=env.action_space,
+                  hidden=64, lr=3e-4, gamma=0.99, clip_coef=0.2,
+                  entropy_coef=0.01, value_coef=0.5, seed=70,
+                  batch_size = 64, ppo_epochs = 4, lam = 0.95)
+    
+    """
+    # Initialize CNN with a dummy observation (to get correct input shape)
+    obs, _ = env.reset()
+    dummy_obs_space = Box(low=0.0, high=1.0, shape=preprocess(obs).shape)
+
+    agent = Agent(obs_space=dummy_obs_space, action_space=env.action_space,
+                  hidden=64, lr=3e-4, gamma=0.99, clip_coef=0.2,
+                  entropy_coef=0.01, value_coef=0.5, seed=70,
+                  batch_size = 64, ppo_epochs = 4, lam = 0.95)
+    """
+    # Stats for Return-Based Scaling only
+    # === Return-Based Scaling stats ===
+    r_mean, r_var = 0.0, 1e-8
+    g2_mean = 1.0
+
+    agent.r_var = r_var
+    agent.g2_mean = g2_mean
+    """
+
+    try:
+        obs, info = env.reset(seed=42)
+        state = preprocess(obs)
+
+        loss_history = []
+        reward_history = []
+
+        for update in range(1, batches + 1):
+            for t in range(steps):
+                action, logp, value = agent.choose_action(state)
+                next_obs, reward, terminated, truncated, info = env.step(action)
+                done = terminated or truncated
+                next_state = preprocess(next_obs)
+
+                agent.remember(state, action, reward, done, logp, value, next_state)
+
+                ep_return += reward
+                state = next_state
+
+                if done:
+                    episode += 1
+                    total_return += ep_return
+                    print(f"Episode {episode} return: {ep_return:.2f}")
+                    ep_return = 0
+                    obs, info = env.reset()
+                    state = preprocess(obs)
+
+            # Using reward gradient clipping 
+            avg_loss = agent.update_reward_gradient_clipping()
+
+            # Vanilla PPO (no normalization)
+            #avg_loss = agent.vanilla_ppo_update()
+            loss_history.append(avg_loss)
+
+            avg_ret = (total_return / episode) if episode else 0
+            reward_history.append(avg_ret)
+            print(f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}, avg_loss={avg_loss:.4f}")
+
+        fig = plt.figure(figsize=(12, 8))
+        
+        """
+        # Plot for Return-Based Scaling only
+        ax1 = plt.subplot(220)
+        ax1.plot(agent.sigma_history, label="Return σ")
+        ax1.set_xlabel("PPO Update")
+        ax1.set_ylabel("σ (Return Std)")
+        """
+
+        ax2 = plt.subplot(221)
+        ax2.plot(loss_history, label="Avg Loss")
+        ax2.set_ylabel("Average PPO Loss")
+        ax2.set_xlabel("PPO Update")
+
+        ax3 = plt.subplot(222)
+        ax3.plot(reward_history, label="Reward")
+        ax3.set_ylabel("Reward")
+        ax3.set_xlabel("PPO Update")
+
+        # Details about value loss and policy loss 
+        ax4 = plt.subplot(223)
+        ax4.plot(agent.policy_loss_history, label="Policy Loss", alpha=0.7)
+        ax4.set_ylabel("Policy Loss")
+        ax4.set_xlabel("Training Step")
+        ax4.legend()
+        
+        ax5 = plt.subplot(224)
+        ax5.plot(agent.value_loss_history, label="Value Loss", alpha=0.7)
+        ax5.set_ylabel("Value Loss")
+        ax5.set_xlabel("Training Step")
+        ax5.legend()
+
+        fig.suptitle("PPO Training Stability")
+        fig.tight_layout()
+        plt.show()
+
+
+
+
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        return 1
+    finally:
+        avg = total_return / episode if episode else 0
+        print(f"\nEpisodes: {episode}, Avg return: {avg:.3f}")
+        env.close()
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())