Add new method of reward clipping

CNN_PPO/ppo_template_meanstd_clipping.py

@@ -0,0 +1,168 @@
+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
+import numpy as np
+
+
+def preprocess(obs):
+    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
+    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
+    return np.expand_dims(obs, axis=0).astype(np.float32) / 255.0
+
+
+def main() -> int:
+    env = gym.make("ALE/Pacman-v5")
+
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 2000
+    batches = 100
+
+    print("Observation space:", env.observation_space)
+    print("Action space:", env.action_space)
+
+    # Initialize CNN
+    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
+
+    # === YOUR NEW REWARD CLIPPING SYSTEM ===
+    alpha = np.random.uniform(0, 2)
+    print(f"\n[INFO] α sampled = {alpha:.3f}\n")
+
+    reward_batch = []          # stores total rewards for 5 episodes
+    clip_low, clip_high = None, None
+    EPISODES_PER_BATCH = 5
+
+    try:
+        obs, info = env.reset()
+        state = preprocess(obs)
+
+        loss_history = []
+        reward_history = []
+
+        # === PPO outer updates ===
+        for update in range(1, batches + 1):
+
+            reward_batch.clear()
+
+            # === Collect 5 full episodes ===
+            for ep in range(EPISODES_PER_BATCH):
+                ep_rewards_raw = []
+                done = False
+
+                while not done:
+                    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)
+
+                    # === APPLY REWARD CLIPPING TO RAW REWARD IF READY ===
+                    # if clip_low is not None:
+                    #     reward = np.clip(reward, clip_low, clip_high)
+
+                    agent.remember(state, action, reward, done, logp, value, next_state)
+
+                    ep_return += reward
+                    ep_rewards_raw.append(reward)
+                    # print("raw reward:", reward)
+                    state = next_state
+
+                    if done:
+                        # episode completed
+                        episode += 1
+                        total_return += ep_return
+
+                        reward_sum = sum(ep_rewards_raw)
+                        if clip_low is not None:
+                            reward_sum_clipped = np.clip(reward_sum, clip_low, clip_high)
+                            reward_batch.append(reward_sum_clipped)
+                            print(f"Episode {episode} | Reward (clipped): {reward_sum_clipped:.2f}")
+                        else:
+                            reward_batch.append(reward_sum)
+                            print(f"Episode {episode} | Reward (clipped): {reward_sum:.2f}")
+
+
+                        ep_return = 0
+                        obs, info = env.reset()
+                        state = preprocess(obs)
+
+            # === After every 5 episodes → compute clipping range ===
+            mu = np.mean(reward_batch)
+            sigma = np.std(reward_batch) + 1e-8
+
+            # clip_low = 0 # When raw reward is clipped
+            clip_low = mu - sigma * alpha # When sum of raw reward is clipped
+            clip_high = mu + alpha * sigma
+
+            print(f"[UPDATE {update}] New Reward Clip Range: [{clip_low:.2f},  {clip_high:.2f}]")
+
+            # === PPO UPDATE ===
+            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}")
+
+        # === PLOTS ===
+        fig = plt.figure(figsize=(12, 8))
+
+        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")
+
+        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())