A2c version on HF

A2C_combo/a2c_helpers.py

@@ -0,0 +1,417 @@
+import numpy as np
+import torch as T
+import torch.nn as nn
+import torch.optim as optim
+from torch.distributions import Categorical
+
+
+class Agent:
+    def __init__(
+            self,
+            obs_space,
+            action_space,
+            hidden,
+            gamma,
+            lr,
+            value_coef,
+            entropy_coef,
+            seed,
+            lam
+
+    ):
+        EPSILON = 1e-8
+
+        # Initialize seed for reproducibility
+        if seed is not None:
+            np.random.seed(seed)
+            T.manual_seed(seed)
+
+        # Use GPU if available
+        self.device = T.device('cuda:0' if T.cuda.is_available() else 'cpu')
+        self.action_dim = int(getattr(action_space, "n", action_space))
+
+        # Initialize the policy and the critic networks
+        self.policy = Policy(obs_space.shape, self.action_dim, hidden).to(self.device)
+        self.critic = Critic(obs_space.shape, hidden).to(self.device)
+
+        # Set optimizer for policy and critic networks
+        self.opt = optim.Adam(
+            list(self.policy.parameters()) + list(self.critic.parameters()),
+            lr=lr
+        )
+
+        self.gamma = gamma
+        self.value_coef = value_coef
+        self.entropy_coef = entropy_coef
+        self.sigma_history = []
+        self.loss_history = []
+        self.policy_loss_history = []
+        self.value_loss_history = []
+        self.entropy_history = []
+        self.lam = lam
+        self.EPSILON = EPSILON
+        self.observeNorm = ObservationNorm()
+        self.advantageNorm = AdvantageNorm()
+        self.returnNorm = ReturnNorm()
+
+        self.memory = Memory()
+
+    # Function to choose action based on current policy
+    # Returns: action, log probabilitiy, value of the state
+    def choose_action(self, observation):
+        state = T.as_tensor(observation, dtype=T.float32, device=self.device)
+        with T.no_grad():
+            dist = self.policy.next_action(state)
+            action = dist.sample()
+            value = self.critic.evaluated_state(state)
+        return int(action.item()), float(value.item())
+
+    # Store reward, state, action in memory
+    def remember(self, state, action, reward, done, value, next_state):
+        with T.no_grad():
+            ns = T.as_tensor(next_state, dtype=T.float32, device=self.device)
+            next_value = self.critic.evaluated_state(ns).item()
+        self.memory.store(state, action, reward, done, value, next_value)
+
+    def _prepare_batch_data(self):
+        """Convert memory to tensors."""
+        states = T.as_tensor(np.array(self.memory.states), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(self.memory.actions, dtype=T.long, device=self.device)
+        rewards = T.as_tensor(self.memory.rewards, dtype=T.float32, device=self.device)
+        dones = T.as_tensor(self.memory.dones, dtype=T.float32, device=self.device)
+        values = T.as_tensor(self.memory.values, dtype=T.float32, device=self.device)
+        return states, actions, rewards, dones, values
+
+    def _compute_gae(self, rewards, values, dones):
+        """Compute Generalized Advantage Estimation."""
+        with T.no_grad():
+            next_values = T.cat([values[1:], values[-1:].clone()])
+            deltas = rewards + self.gamma * next_values * (1 - dones) - values
+
+            adv = T.zeros_like(rewards)
+            gae = 0.0
+            for t in reversed(range(len(rewards))):
+                gae = deltas[t] + self.gamma * self.lam * (1 - dones[t]) * gae
+                adv[t] = gae
+
+            returns = adv + values
+            return adv, returns
+
+    def _compute_a2c_loss(self, states, actions, returns, advantages):
+        """Compute A2C loss components."""
+        dist = self.policy.next_action(states)
+        new_logp = dist.log_prob(actions)
+        entropy = dist.entropy().mean()
+
+        # Simple policy gradient (no clipping)
+        policy_loss = -(new_logp * advantages).mean()
+
+        # Critic loss
+        value_pred = self.critic.evaluated_state(states)
+        value_loss = 0.5 * (returns - value_pred).pow(2).mean()
+
+        # Total loss
+        total_loss = (
+                policy_loss +
+                self.value_coef * value_loss -
+                self.entropy_coef * entropy
+        )
+
+        return total_loss, policy_loss, value_loss
+
+    def _a2c_update(self, states, actions, returns, adv, use_grad_clip=False):
+        """Run single A2C update (no multiple epochs)."""
+        total_loss, policy_loss, value_loss = self._compute_a2c_loss(
+            states, actions, returns, adv
+        )
+
+        self.policy_loss_history.append(policy_loss.item())
+        self.value_loss_history.append(value_loss.item())
+
+        self.opt.zero_grad(set_to_none=True)
+        total_loss.backward()
+
+        if use_grad_clip:
+            T.nn.utils.clip_grad_norm_(
+                list(self.policy.parameters()) + list(self.critic.parameters()),
+                0.5
+            )
+
+        self.opt.step()
+
+        return total_loss.item()
+
+    def vanilla_a2c_update(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states, actions, rewards, dones, values = self._prepare_batch_data()
+        adv, returns = self._compute_gae(rewards, values, dones)
+
+        with T.no_grad():
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + self.EPSILON)
+
+        avg_loss = self._a2c_update(states, actions, returns, adv)  # changed
+        self.memory.clear()
+        return avg_loss
+
+    def update_rbs(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states, actions, rewards, dones, values = self._prepare_batch_data()
+        adv, returns = self._compute_gae(rewards, values, dones)
+
+        with T.no_grad():
+            sigma_t = returns.std(unbiased=False) + 1e-8
+            returns = returns / sigma_t
+            self.sigma_history.append(sigma_t.item())
+            adv = adv / sigma_t
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        avg_loss = self._a2c_update(states, actions, returns, adv)  # changed
+        self.memory.clear()
+        return avg_loss
+
+    def update_gradient_clipping(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states, actions, rewards, dones, values = self._prepare_batch_data()
+        adv, returns = self._compute_gae(rewards, values, dones)
+
+        with T.no_grad():
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        avg_loss = self._a2c_update(states, actions, returns, adv, use_grad_clip=True)  # changed
+        self.memory.clear()
+        return avg_loss
+
+    def update_obs_norm(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states, actions, rewards, dones, values = self._prepare_batch_data()
+        adv, returns = self._compute_gae(rewards, values, dones)
+
+        with T.no_grad():
+            # --- observation normalization ---
+            states = self.observeNorm.normalize(states)
+            # Advantage normalization
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        avg_loss = self._a2c_update(states, actions, returns, adv)
+        self.memory.clear()
+        return avg_loss
+
+    def update_adv_norm(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states, actions, rewards, dones, values = self._prepare_batch_data()
+        adv, returns = self._compute_gae(rewards, values, dones)
+
+        with T.no_grad():
+            # --- Advantage normalization ---
+            adv = self.advantageNorm.normalize(adv)
+
+        avg_loss = self._a2c_update(states, actions, returns, adv)
+        self.memory.clear()
+        return avg_loss
+
+    def update_return_norm(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states = T.as_tensor(np.array(self.memory.states), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(self.memory.actions, dtype=T.long, device=self.device)
+        rewards = T.as_tensor(self.memory.rewards, dtype=T.float32, device=self.device)
+        dones = T.as_tensor(self.memory.dones, dtype=T.float32, device=self.device)
+        values = T.as_tensor(self.memory.values, dtype=T.float32, device=self.device)
+
+        with T.no_grad():
+            next_values = T.cat([values[1:], values[-1:].clone()])
+            returns = rewards + self.gamma * next_values * (1 - dones)
+            adv = returns - values
+            returns = self.returnNorm.normalize(returns)
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        dist = self.policy.next_action(states)
+        log_probs = dist.log_prob(actions)
+        entropy = dist.entropy().mean()
+
+        policy_loss = -(log_probs * adv).mean()
+
+        value_pred = self.critic.evaluated_state(states)
+        value_loss = 0.5 * (returns - value_pred).pow(2).mean()
+
+        total_loss = (
+                policy_loss +
+                self.value_coef * value_loss -
+                self.entropy_coef * entropy
+        )
+
+        self.opt.zero_grad(set_to_none=True)
+        total_loss.backward()
+        self.opt.step()
+
+        avg_loss = self._a2c_update(states, actions, returns, adv)
+        self.memory.clear()
+        return avg_loss
+
+    def update_reward_norm(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        states = T.as_tensor(np.array(self.memory.states), dtype=T.float32, device=self.device)
+        actions = T.as_tensor(self.memory.actions, dtype=T.long, device=self.device)
+        rewards = T.as_tensor(self.memory.rewards, dtype=T.float32, device=self.device)
+        dones = T.as_tensor(self.memory.dones, dtype=T.float32, device=self.device)
+        values = T.as_tensor(self.memory.values, dtype=T.float32, device=self.device)
+
+        rewards = (rewards - rewards.mean()) / (rewards.std(unbiased=False) + 1e-8)
+
+        with T.no_grad():
+            next_values = T.cat([values[1:], values[-1:].clone()])
+
+            returns = rewards + self.gamma * next_values * (1 - dones)
+
+            adv = returns - values
+
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        dist = self.policy.next_action(states)
+        log_probs = dist.log_prob(actions)
+        entropy = dist.entropy().mean()
+
+        # Actor Loss: - log_prob * Advantage
+        policy_loss = -(log_probs * adv).mean()
+
+        value_pred = self.critic.evaluated_state(states)
+        value_loss = 0.5 * (returns - value_pred).pow(2).mean()
+
+        total_loss = (
+                policy_loss +
+                self.value_coef * value_loss -
+                self.entropy_coef * entropy
+        )
+
+        self.opt.zero_grad(set_to_none=True)
+        total_loss.backward()
+        self.opt.step()
+
+        avg_loss = self._a2c_update(states, actions, returns, adv)
+        self.memory.clear()
+        return avg_loss
+
+# Policy network (CNN)
+class Policy(nn.Module):
+    def __init__(self, obs_shape: tuple, action_dim: int, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        # Suggested architecture for Atari: https://arxiv.org/pdf/1312.5602
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten(),
+            nn.Linear(32 * 9 * 9, 256),  # 2592 → 256
+            nn.ReLU(),
+        )
+
+        # Final output layer: one logit per action
+        self.net = nn.Linear(256, action_dim)
+
+    def next_action(self, state: T.Tensor) -> Categorical:
+        # state shape should be (B, C, H, W)
+        if state.dim() == 3:
+            state = state.unsqueeze(0)
+
+        cnn_out = self.cnn(state)  # [B, 256]
+        logits = self.net(cnn_out)  # [B, action_dim]
+        return Categorical(logits=logits)
+
+
+# Critic network (CNN)
+class Critic(nn.Module):
+    def __init__(self, obs_shape: tuple, hidden: int):
+        super().__init__()
+        c, h, w = obs_shape
+        # Suggested architecture for Atari: https://arxiv.org/pdf/1312.5602
+        self.cnn = nn.Sequential(
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
+            nn.ReLU(),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
+            nn.ReLU(),
+            nn.Flatten()
+        )
+
+        with T.no_grad():
+            cnn_output_dim = self.cnn(T.zeros(1, c, h, w)).shape[1]
+
+        self.net = nn.Sequential(
+            nn.Linear(cnn_output_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, 1)
+        )
+
+    def evaluated_state(self, x: T.Tensor) -> T.Tensor:
+        if x.dim() == 3:
+            x = x.unsqueeze(0)
+        cnn_out = self.cnn(x)
+        return self.net(cnn_out).squeeze(-1)
+
+
+class Memory():
+    def __init__(self):
+        self.states = []
+        self.actions = []
+        self.rewards = []
+        self.dones = []
+        self.values = []
+        self.next_values = []
+
+    def store(self, state, action, reward, done, value, next_value):
+        self.states.append(np.asarray(state, dtype=np.float32))
+        self.actions.append(int(action))
+        self.rewards.append(float(reward))
+        self.dones.append(float(done))
+        self.values.append(float(value))
+        self.next_values.append(float(next_value))
+
+    def clear(self):
+        self.states = []
+        self.actions = []
+        self.rewards = []
+        self.dones = []
+        self.values = []
+        self.next_values = []
+
+
+class ObservationNorm:
+
+    def normalize(self, x):
+        return (x - x.mean()) / (x.std(unbiased=False) + 1e-8)  # We add epsilon to make sure that we don't
+        # divide through zero.
+
+
+class AdvantageNorm:
+    '''
+    This class implements the Advantage Normalization. The purpose is to normalize either across batches or
+    only within the same batch.
+    '''
+
+    def normalize(self, x):
+        return (x - x.mean()) / (x.std(unbiased=False) + 1e-8)  # We add epsilon to make sure that we don't
+        # divide through zero.
+
+
+class ReturnNorm:
+    '''
+    This class implements the Return Normalization. The purpose is to normalize either across batches or
+    only within the same batch.
+    '''
+
+    def normalize(self, x):
+        return (x - x.mean()) / (x.std(unbiased=False) + 1e-8)

A2C_combo/a2c_main.py

@@ -0,0 +1,348 @@
+import argparse
+import gymnasium as gym
+import sys
+import matplotlib.pyplot as plt
+import ale_py
+from a2c_helpers import *
+from gymnasium.spaces import Box
+import cv2
+import logging
+import numpy as np
+import pandas as pd
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
+logger = logging.getLogger(__name__)
+
+
+# Preprocess environment
+def preprocess(obs):
+    # Convert to grayscale
+    obs = cv2.cvtColor(obs, cv2.COLOR_RGB2GRAY)
+    # Resize
+    obs = cv2.resize(obs, (84, 84), interpolation=cv2.INTER_AREA)
+
+    return np.expand_dims(obs, axis=0).astype(np.float32) / 255.0
+
+
+def df_ops(lst_df, seeds):
+    for df in lst_df:
+        seed_data = df[seeds]
+        df['Avg'] = seed_data.mean(axis=1)
+        df['High'] = seed_data.max(axis=1)
+        df['Low'] = seed_data.min(axis=1)
+
+    return lst_df
+
+
+# Main loop
+def main() -> int:
+    # Initialize variables
+    batches = 1000
+    steps = 5
+    clip_interval = 2
+    seeds = [10, 20, 30, 40, 50]
+    ep_per_batch = 5
+
+    # batches = 5
+    # steps = 5
+    # clip_interval = 2
+    # seeds = [10, 20]
+    # ep_per_batch = 2
+    # Arguments 
+    """
+    usage examples:
+    python3 a2c_main.py --method vanilla
+
+    python3 a2c_main.py --method grad_clip
+
+    python3 a2c_main.py --method rbs 
+    """
+    parser = argparse.ArgumentParser(description='A2C Training')
+
+    parser.add_argument('--method', type=str, choices=['vanilla', 'reward_clip', 'rbs', 'grad_clip',
+                                                       'obs_norm', 'adv_norm', 'return_norm', 'reward_norm'],
+                        default='vanilla', help='A2C update method')
+    parser.add_argument('--env', type=str, default='ALE/Pacman-v5',
+                        help='Gym environment name (e.g., ALE/Pacman-v5, ALE/SpaceInvaders-v5)')
+    parser.add_argument('--render', action='store_true', help='Enable rendering')
+    parser.add_argument('--clip_window', type=int, default=clip_interval,
+                        help='Number of batches to collect rewards for clipping range update')
+
+    args = parser.parse_args()
+
+    # Set up environment
+    if args.render:
+        env = gym.make(args.env, render_mode="human")
+    else:
+        env = gym.make(args.env)
+
+    logger.info(f"Observation space: {env.observation_space}")
+    logger.info(f"Action space: {env.action_space}")
+    logger.info(f'Method: {args.method}')
+
+    # 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)
+
+    # Initialize A2C agent
+    agent = Agent(obs_space=dummy_obs_space, action_space=env.action_space,
+                  hidden=64, lr=0.00001, gamma=0.997,
+                  entropy_coef=0.01, value_coef=0.5, seed=70,lam=0.95)
+
+    # === Return-Based Scaling stats (for RBS method) ===
+    r_mean, r_var = 0.0, 1e-8
+    g2_mean = 1.0
+    agent.r_var = r_var
+    agent.g2_mean = g2_mean
+
+    all_reward_histories = pd.DataFrame(columns=[i for i in seeds], index=[i for i in range(1, batches + 1)])
+    all_loss_histories = pd.DataFrame(columns=[i for i in seeds], index=[i for i in range(1, batches + 1)])
+    all_policy_loss = pd.DataFrame(columns=[i for i in seeds])
+    all_value_loss = pd.DataFrame(columns=[i for i in seeds])
+
+    step = 0
+    # Main update loop
+    try:
+        for seed in seeds:
+            obs, info = env.reset(seed=seed)
+            state = preprocess(obs)
+
+            loss_history = []
+            reward_history = []
+            episode = 0
+            total_return = 0
+
+            """ Update loop: Gradient, Reward Normalization """
+            if args.method == 'reward_clip':
+                alpha = np.random.uniform(1, 2)
+                logger.info(f"α sampled = {alpha:.3f} seed = {seed}")
+
+                clip_low, clip_high = None, None
+                ep_reward_history = []
+
+                obs, info = env.reset()
+                state = preprocess(obs)
+
+                for update in range(1, batches + 1):
+
+                    batch_episode_returns = []  # used for μ, σ
+
+                    for _ in range(ep_per_batch):
+                        ep_rewards = []
+                        done = False
+
+                        while not done:
+                            action, value = agent.choose_action(state)
+                            next_obs, reward, terminated, truncated, info = env.step(action)
+                            done = terminated or truncated
+                            next_state = preprocess(next_obs)
+
+                            ep_rewards.append(reward)
+
+                            agent.remember(state, action, reward, done, value, next_state)
+
+                            state = next_state
+
+                            if done:
+                                ep_return = sum(ep_rewards)
+                                if clip_low is not None:
+                                    clipped_return = np.clip(ep_return, clip_low, clip_high)
+                                else:
+                                    clipped_return = ep_return
+                                ep_reward_history.append(clipped_return)
+                                batch_episode_returns.append(clipped_return)
+
+                                episode += 1
+                                total_return += clipped_return
+
+                                logger.info(f"Episode {episode} return: {clipped_return:.2f}")
+
+                                obs, info = env.reset()
+                                state = preprocess(obs)
+
+                    # === Compute clipping bounds using Code 1 logic ===
+                    mu = np.mean(batch_episode_returns)
+                    sigma = np.std(batch_episode_returns) + 1e-8 if np.std(batch_episode_returns)!=0 else 1
+
+                    clip_low = mu - alpha * sigma
+                    clip_high = mu + alpha * sigma
+
+                    logger.info(
+                        f"[UPDATE {update}] New Reward Clip Range: "
+                        f"[{clip_low:.4f}, {clip_high:.4f}]"
+                    )
+
+                    # === A2C UPDATE ===
+                    avg_loss = agent.vanilla_a2c_update()
+                    loss_history.append(avg_loss)
+
+                    avg_ret = np.mean(batch_episode_returns)
+                    reward_history.append(avg_ret)
+
+                    logger.info(
+                        f"Update {update}: batch_mean={avg_ret:.4f}, "
+                        f"batch_std={np.std(batch_episode_returns):.4f}, "
+                        f"episodes={episode}, avg_loss={avg_loss:.4f}"
+                    )
+
+                """ Update loop: Other Normalization Methods """
+            else:
+                for update in range(1, batches + 1):
+                    batch_episode_rewards = []
+                    ep_per_batch = 5
+
+                    for _ in range(ep_per_batch):
+                        ep_rewards = []
+
+                        done = False
+
+                        while not done:
+                            action, value = agent.choose_action(state)
+                            next_obs, reward, terminated, truncated, info = env.step(action)
+                            done = terminated or truncated
+                            next_state = preprocess(next_obs)
+
+                            ep_rewards.append(reward)
+                            agent.remember(state, action, reward, done, value, next_state)
+
+                            state = next_state
+
+                            if done:
+                                ep_return = sum(ep_rewards)
+                                episode += 1
+                                total_return += ep_return
+                                batch_episode_rewards.append(ep_return)
+                                logger.info(f"Episode {episode} return: {ep_return:.2f}")
+
+                                obs, info = env.reset()
+                                state = preprocess(obs)
+
+                    # Choose normalization method
+                    if args.method == 'vanilla':
+                        avg_loss = agent.vanilla_a2c_update()
+                    elif args.method == 'grad_clip':
+                        avg_loss = agent.update_gradient_clipping()
+                    elif args.method == 'obs_norm':
+                        avg_loss = agent.update_obs_norm()
+                    elif args.method == 'adv_norm':
+                        avg_loss = agent.update_adv_norm()
+                    elif args.method == 'reward_norm':
+                        avg_loss = agent.update_reward_norm()
+                    else:  # rbs
+                        avg_loss = agent.update_rbs()
+
+                    loss_history.append(avg_loss)
+
+                    avg_ret = (total_return / episode) if episode else 0
+                    reward_history.append(avg_ret)
+                    logger.info(
+                        f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}, avg_loss={avg_loss:.4f}")
+
+            all_reward_histories[seed] = reward_history
+            all_loss_histories[seed] = loss_history
+            all_value_loss[seed] = agent.value_loss_history[step:step+batches]
+            all_policy_loss[seed] = agent.policy_loss_history[step:step+batches]
+
+            step += batches
+
+        [all_reward_histories, all_loss_histories, all_value_loss, all_policy_loss] = df_ops([all_reward_histories,
+                                                                                              all_loss_histories,
+                                                                                              all_value_loss,
+                                                                                              all_policy_loss],
+                                                                                             seeds)
+        # [all_reward_histories, all_loss_histories] = df_ops([all_reward_histories, all_loss_histories], seeds)
+
+        # all_policy_loss.to_csv(args.method + '_policy_loss.csv')
+        # all_value_loss.to_csv(args.method + '_value_loss.csv')
+        # all_reward_histories.to_csv(args.method + '_a2c_reward_history.csv')
+        # all_loss_histories.to_csv(args.method + '_a2c_loss_history.csv')
+
+        fig = plt.figure(figsize=(15, 10))
+
+        # --- Subplot 1: Average PPO Loss ---
+        ax2 = plt.subplot(221)
+        # Plot the shaded High-Low Range
+        ax2.fill_between(
+            all_loss_histories.index,
+            all_loss_histories['Low'],
+            all_loss_histories['High'],
+            color='#A8DADC',  # Light blue for aesthetic shading
+            alpha=0.5,
+            label="High-Low Range"
+        )
+        # Plot the Average Line
+        ax2.plot(all_loss_histories['Avg'], label="Avg Loss", color='#1D3557', linewidth=2)
+        ax2.set_ylabel("Average PPO Loss")
+        ax2.set_xlabel("PPO Update")
+        ax2.legend()
+
+        # --- Subplot 2: Reward ---
+        ax3 = plt.subplot(222)
+        # Plot the shaded High-Low Range
+        ax3.fill_between(
+            all_reward_histories.index,
+            all_reward_histories['Low'],
+            all_reward_histories['High'],
+            color='#FEDCC8',  # Light orange/peach
+            alpha=0.5,
+            label="High-Low Range"
+        )
+        # Plot the Average Line
+        ax3.plot(all_reward_histories['Avg'], label="Avg Reward", color='#E63946', linewidth=2)
+        ax3.set_ylabel("Average Reward")
+        ax3.set_xlabel("PPO Update")
+        ax3.legend()
+
+        # --- Subplot 3: Policy Loss ---
+        ax4 = plt.subplot(223)
+        # Plot the shaded High-Low Range
+        ax4.fill_between(
+            all_policy_loss.index,
+            all_policy_loss['Low'],
+            all_policy_loss['High'],
+            color='#B0E0A0',  # Light green
+            alpha=0.5,
+            label="High-Low Range"
+        )
+        # Plot the Average Line
+        ax4.plot(all_policy_loss['Avg'], label="Policy Loss", color='#38B000', linewidth=2)
+        ax4.set_ylabel("Average Policy Loss")
+        ax4.set_xlabel("PPO Update")
+        ax4.legend()
+
+        # --- Subplot 4: Value Loss ---
+        ax5 = plt.subplot(224)
+        # Plot the shaded High-Low Range
+        ax5.fill_between(
+            all_value_loss.index,
+            all_value_loss['Low'],
+            all_value_loss['High'],
+            color='#D7BDE2',  # Light purple
+            alpha=0.5,
+            label="High-Low Range"
+        )
+        # Plot the Average Line
+        ax5.plot(all_value_loss['Avg'], label="Value Loss", color='#8E44AD', linewidth=2)
+        ax5.set_ylabel("Average Value Loss")
+        ax5.set_xlabel("PPO Update")
+        ax5.legend()
+
+        # --- Figure Settings ---
+        fig.suptitle(f"PPO Training Stability - {args.method}", fontsize=16, fontweight='bold')
+        # fig.tight_layout()  # Adjust layout to make room for suptitle
+        plt.show()
+
+    except Exception as e:
+        logger.error(f"Error: {e}", exc_info=True)
+        return 1
+    finally:
+        avg = total_return / episode if episode else 0
+        logger.info(f"\nEpisodes: {episode}, Avg return: {avg:.3f}")
+        env.close()
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())