Create test.py

test.py

@@ -0,0 +1,403 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.distributions import Categorical
+import numpy as np
+import gymnasium as gym
+import matplotlib.pyplot as plt
+from collections import deque
+
+# Set random seeds for reproducibility
+torch.manual_seed(42)
+np.random.seed(42)
+
+# Check if GPU is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+
+# ==================== Policy Networks ====================
+
+class CartPolePolicy(nn.Module):
+    """Policy network for CartPole environment"""
+    def __init__(self, state_dim, action_dim, hidden_dim=128):
+        super(CartPolePolicy, self).__init__()
+        self.fc1 = nn.Linear(state_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, action_dim)
+    
+    def forward(self, x):
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.softmax(x, dim=-1)
+
+
+class PongPolicy(nn.Module):
+    """Policy network for Pong with CNN architecture"""
+    def __init__(self, action_dim=2):
+        super(PongPolicy, self).__init__()
+        # CNN layers for processing 80x80 images
+        self.conv1 = nn.Conv2d(1, 16, kernel_size=8, stride=4)
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=4, stride=2)
+        
+        # Calculate size after convolutions: 80 -> 19 -> 8
+        self.fc1 = nn.Linear(32 * 8 * 8, 256)
+        self.fc2 = nn.Linear(256, action_dim)
+    
+    def forward(self, x):
+        # x shape: (batch, 80, 80) -> add channel dimension
+        if len(x.shape) == 2:
+            x = x.unsqueeze(0).unsqueeze(0)
+        elif len(x.shape) == 3:
+            x = x.unsqueeze(1)
+        
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.softmax(x, dim=-1)
+
+
+# ==================== Helper Functions ====================
+
+def preprocess(image):
+    """Prepro 210x160x3 uint8 frame into 6400 (80x80) 2D float array"""
+    image = image[35:195]  # crop
+    image = image[::2, ::2, 0]  # downsample by factor of 2
+    image[image == 144] = 0  # erase background (background type 1)
+    image[image == 109] = 0  # erase background (background type 2)
+    image[image != 0] = 1  # everything else (paddles, ball) just set to 1
+    return np.reshape(image.astype(float).ravel(), [80, 80])
+
+
+def compute_returns(rewards, gamma):
+    """Compute discounted returns for each timestep"""
+    returns = []
+    R = 0
+    for r in reversed(rewards):
+        R = r + gamma * R
+        returns.insert(0, R)
+    returns = torch.tensor(returns, dtype=torch.float32).to(device)
+    # Normalize returns for more stable training
+    if len(returns) > 1:
+        returns = (returns - returns.mean()) / (returns.std() + 1e-8)
+    return returns
+
+
+def moving_average(data, window_size):
+    """Compute moving average"""
+    if len(data) < window_size:
+        return np.array([np.mean(data[:i+1]) for i in range(len(data))])
+    
+    moving_avg = []
+    for i in range(len(data)):
+        if i < window_size:
+            moving_avg.append(np.mean(data[:i+1]))
+        else:
+            moving_avg.append(np.mean(data[i-window_size+1:i+1]))
+    return np.array(moving_avg)
+
+
+# ==================== Policy Gradient Algorithm ====================
+
+def train_policy_gradient(env_name, policy, optimizer, gamma, num_episodes, 
+                         max_steps=None, is_pong=False, action_map=None):
+    """
+    Train policy using REINFORCE algorithm
+    
+    Args:
+        env_name: Name of the gym environment
+        policy: Policy network
+        optimizer: PyTorch optimizer
+        gamma: Discount factor
+        num_episodes: Number of training episodes
+        max_steps: Maximum steps per episode (None for default)
+        is_pong: Whether this is Pong environment
+        action_map: Mapping from policy action to env action (for Pong)
+    """
+    env = gym.make(env_name)
+    episode_rewards = []
+    
+    for episode in range(num_episodes):
+        state, _ = env.reset()
+        
+        # Preprocess state for Pong
+        if is_pong:
+            state = preprocess(state)
+            prev_frame = None  # Track previous frame for motion
+        
+        log_probs = []
+        rewards = []
+        
+        done = False
+        step = 0
+        
+        while not done:
+            # For Pong, use frame difference (motion signal)
+            if is_pong:
+                cur_frame = state
+                if prev_frame is not None:
+                    state_input = cur_frame - prev_frame
+                else:
+                    state_input = np.zeros_like(cur_frame, dtype=np.float32)
+                prev_frame = cur_frame
+                state_tensor = torch.FloatTensor(state_input).to(device)
+            else:
+                # Convert state to tensor
+                state_tensor = torch.FloatTensor(state).to(device)
+            
+            # Get action probabilities
+            action_probs = policy(state_tensor)
+            
+            # Sample action from the distribution
+            dist = Categorical(action_probs)
+            action = dist.sample()
+            log_prob = dist.log_prob(action)
+            
+            # Map action for Pong (0,1 -> 2,3)
+            if is_pong:
+                env_action = action_map[action.item()]
+            else:
+                env_action = action.item()
+            
+            # Take action in environment
+            next_state, reward, terminated, truncated, _ = env.step(env_action)
+            done = terminated or truncated
+            
+            # Preprocess next state for Pong
+            if is_pong:
+                next_state = preprocess(next_state)
+            
+            # Store log probability and reward
+            log_probs.append(log_prob)
+            rewards.append(reward)
+            
+            state = next_state
+            step += 1
+            
+            if max_steps and step >= max_steps:
+                break
+        
+        # Compute returns
+        returns = compute_returns(rewards, gamma)
+        
+        # Compute policy gradient loss
+        policy_loss = []
+        for log_prob, R in zip(log_probs, returns):
+            policy_loss.append(-log_prob * R)
+        
+        # Optimize policy
+        optimizer.zero_grad()
+        policy_loss = torch.stack(policy_loss).sum()
+        policy_loss.backward()
+        optimizer.step()
+        
+        # Record episode reward
+        episode_reward = sum(rewards)
+        episode_rewards.append(episode_reward)
+        
+        # Print progress
+        if (episode + 1) % 100 == 0:
+            avg_reward = np.mean(episode_rewards[-100:])
+            print(f"Episode {episode + 1}/{num_episodes}, "
+                  f"Avg Reward (last 100): {avg_reward:.2f}")
+    
+    env.close()
+    return episode_rewards
+
+
+def evaluate_policy(env_name, policy, num_episodes=500, is_pong=False, action_map=None):
+    """Evaluate trained policy over multiple episodes"""
+    env = gym.make(env_name)
+    eval_rewards = []
+    
+    for episode in range(num_episodes):
+        state, _ = env.reset()
+        
+        if is_pong:
+            state = preprocess(state)
+            prev_frame = None  # Track previous frame for motion
+        
+        episode_reward = 0
+        done = False
+        
+        while not done:
+            # For Pong, use frame difference (motion signal)
+            if is_pong:
+                cur_frame = state
+                if prev_frame is not None:
+                    state_input = cur_frame - prev_frame
+                else:
+                    state_input = np.zeros_like(cur_frame, dtype=np.float32)
+                prev_frame = cur_frame
+                state_tensor = torch.FloatTensor(state_input).to(device)
+            else:
+                state_tensor = torch.FloatTensor(state).to(device)
+            
+            with torch.no_grad():
+                action_probs = policy(state_tensor)
+                action = torch.argmax(action_probs).item()
+            
+            if is_pong:
+                env_action = action_map[action]
+            else:
+                env_action = action
+            
+            next_state, reward, terminated, truncated, _ = env.step(env_action)
+            done = terminated or truncated
+            
+            if is_pong:
+                next_state = preprocess(next_state)
+            
+            episode_reward += reward
+            state = next_state
+        
+        eval_rewards.append(episode_reward)
+        
+        if (episode + 1) % 100 == 0:
+            print(f"Evaluated {episode + 1}/{num_episodes} episodes")
+    
+    env.close()
+    return eval_rewards
+
+
+def plot_results(episode_rewards, eval_rewards, title, save_prefix):
+    """Plot training curve and evaluation histogram"""
+    fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+    
+    # Plot training curve
+    ax1 = axes[0]
+    episodes = np.arange(1, len(episode_rewards) + 1)
+    ma = moving_average(episode_rewards, 100)
+    
+    ax1.plot(episodes, episode_rewards, alpha=0.3, label='Episode Reward')
+    ax1.plot(episodes, ma, linewidth=2, label='Moving Average (100 episodes)')
+    ax1.set_xlabel('Episode')
+    ax1.set_ylabel('Reward')
+    ax1.set_title(f'{title} - Training Curve')
+    ax1.legend()
+    ax1.grid(True, alpha=0.3)
+    
+    # Plot evaluation histogram
+    ax2 = axes[1]
+    mean_reward = np.mean(eval_rewards)
+    std_reward = np.std(eval_rewards)
+    
+    ax2.hist(eval_rewards, bins=30, edgecolor='black', alpha=0.7)
+    ax2.axvline(mean_reward, color='red', linestyle='--', linewidth=2, 
+                label=f'Mean: {mean_reward:.2f}')
+    ax2.set_xlabel('Episode Reward')
+    ax2.set_ylabel('Frequency')
+    ax2.set_title(f'{title} - Evaluation Histogram (500 episodes)\n'
+                  f'Mean: {mean_reward:.2f}, Std: {std_reward:.2f}')
+    ax2.legend()
+    ax2.grid(True, alpha=0.3, axis='y')
+    
+    plt.tight_layout()
+    plt.savefig(f'{save_prefix}_results.png', dpi=150, bbox_inches='tight')
+    plt.show()
+    
+    print(f"\n{title} Evaluation Results:")
+    print(f"Mean Reward: {mean_reward:.2f}")
+    print(f"Std Reward: {std_reward:.2f}")
+
+
+# ==================== Main Training Scripts ====================
+
+def train_cartpole():
+    """Train CartPole-v1"""
+    print("\n" + "="*60)
+    print("Training CartPole-v1")
+    print("="*60 + "\n")
+    
+    # Environment parameters
+    env = gym.make('CartPole-v1')
+    state_dim = env.observation_space.shape[0]
+    action_dim = env.action_space.n
+    env.close()
+    
+    # Hyperparameters
+    gamma = 0.95
+    learning_rate = 0.01
+    num_episodes = 1000
+    
+    # Initialize policy and optimizer
+    policy = CartPolePolicy(state_dim, action_dim).to(device)
+    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)
+    
+    # Train
+    episode_rewards = train_policy_gradient(
+        'CartPole-v1', policy, optimizer, gamma, num_episodes
+    )
+    
+    # Evaluate
+    print("\nEvaluating trained policy...")
+    eval_rewards = evaluate_policy('CartPole-v1', policy, num_episodes=500)
+    
+    # Plot results
+    plot_results(episode_rewards, eval_rewards, 'CartPole-v1', 'cartpole')
+    
+    # Save model
+    torch.save(policy.state_dict(), 'cartpole_policy.pth')
+    print("\nModel saved as 'cartpole_policy.pth'")
+    
+    return policy, episode_rewards, eval_rewards
+
+
+def train_pong():
+    """Train Pong-v5"""
+    print("\n" + "="*60)
+    print("Training Pong-v5")
+    print("="*60 + "\n")
+    
+    # Hyperparameters
+    gamma = 0.99
+    learning_rate = 0.001  # Lower learning rate for Pong
+    num_episodes = 3000  # Pong requires more episodes
+    
+    # Action mapping: policy outputs 0 or 1, map to RIGHT(2) or LEFT(3)
+    action_map = {0: 2, 1: 3}  # 0->RIGHT, 1->LEFT
+    
+    # Initialize policy and optimizer
+    policy = PongPolicy(action_dim=2).to(device)
+    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)
+    
+    print(f"Using learning rate: {learning_rate}")
+    print(f"Action mapping: 0->RIGHT(2), 1->LEFT(3)\n")
+    
+    # Train
+    episode_rewards = train_policy_gradient(
+        'ALE/Pong-v5', policy, optimizer, gamma, num_episodes,
+        is_pong=True, action_map=action_map
+    )
+    
+    # Evaluate
+    print("\nEvaluating trained policy...")
+    eval_rewards = evaluate_policy(
+        'ALE/Pong-v5', policy, num_episodes=500,
+        is_pong=True, action_map=action_map
+    )
+    
+    # Plot results
+    plot_results(episode_rewards, eval_rewards, 'Pong-v5', 'pong')
+    
+    # Save model
+    torch.save(policy.state_dict(), 'pong_policy.pth')
+    print("\nModel saved as 'pong_policy.pth'")
+    
+    return policy, episode_rewards, eval_rewards
+
+
+# ==================== Run Training ====================
+
+if __name__ == "__main__":
+    # Train CartPole
+    cartpole_policy, cartpole_train_rewards, cartpole_eval_rewards = train_cartpole()
+    
+    # Train Pong (this will take longer)
+    print("\n\nNote: Pong training will take significantly longer (may take hours)")
+    print("You may want to reduce num_episodes if just testing the code.\n")
+    
+    # Uncomment the line below to train Pong
+    # pong_policy, pong_train_rewards, pong_eval_rewards = train_pong()