Adrian's first attempt of PPO

.gitignore

@@ -0,0 +1 @@
+.idea/

Adrian/ppo_helpers_v2 (1).py

@@ -0,0 +1,233 @@
+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():
+    # Minimal PPO-Clip agent (single full-batch update per episode, MC returns)
+    def __init__(
+        self,
+        obs_space,
+        action_space,
+        hidden,
+        gamma,
+        clip_coef,
+        lr,
+        value_coef,
+        entropy_coef,
+        seed
+    ):
+        # 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.obs_dim = int(np.prod(getattr(obs_space, "shape", (obs_space,))))
+        self.action_dim = int(getattr(action_space, "n", action_space))
+
+        # Initialize the policy and the critic networks
+        self.policy = Policy(self.obs_dim, self.action_dim, hidden).to(self.device)
+        self.critic = Critic(self.obs_dim, 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.clip = clip_coef
+        self.value_coef = value_coef
+        self.entropy_coef = entropy_coef
+
+        self.memory = Memory()
+
+    def choose_action(self, observation):
+        # Returns: action, log probabilitiy, value of the state
+        state = T.as_tensor(observation, dtype=T.float32, device=self.device).view(-1)
+        with T.no_grad():
+            # Forward function (defined in Policy class)
+            dist = self.policy.next_action(state)
+            action = dist.sample()
+            logp = dist.log_prob(action)
+            value = self.critic.evaluated_state(state)
+        return int(action.item()), float(logp.item()), float(value.item())
+
+
+    def remember(self, state, action, reward, done, log_prob, value, next_state):
+        with T.no_grad():
+            # Pass on next state and have it evaluated by the critic network
+            ns = T.as_tensor(next_state, dtype=T.float32, device=self.device).view(-1)
+            next_value = self.critic.evaluated_state(ns).item()
+        self.memory.store(state, action, reward, done, log_prob, value, next_value)
+
+    """
+    def run_episode(self, env, max_steps: int, render: bool = False):
+        # Runs one episode, updates the policy once at the end
+        self.memory.clear()
+        out = env.reset()
+
+        state = out[0] if isinstance(out, tuple) else out
+
+        ep_return, ep_len = 0, 0
+
+        steps_limit = max_steps if max_steps is not None else float("inf")
+
+        while ep_len < steps_limit:
+            if render and hasattr(env, "render"):
+                env.render()
+
+            action, logp, value = self.choose_action(state)
+            step_out = env.step(action)
+            if len(step_out) == 5:
+                next_state, reward, terminated, truncated, _ = step_out
+                done = terminated or truncated
+            else:
+                next_state, reward, done, _ = step_out
+
+            self.remember(state, action, reward, done, logp, value, next_state)
+
+            ep_return += float(reward)
+            ep_len += 1
+            state = next_state
+            if done:
+                break
+
+        self._update()
+        return ep_return, ep_len
+
+    def run_episodes(self, env, n_episodes: int, max_steps: int, render: bool = False):
+        returns = []
+        for _ in range(n_episodes):
+            ep_ret, _ = self.run_episode(env, max_steps=max_steps, render=render)
+            returns.append(ep_ret)
+        return returns
+    
+    """
+
+    def _update(self):
+        if len(self.memory.states) == 0:
+            return
+
+        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)
+        old_logp = T.as_tensor(self.memory.log_probs, dtype=T.float32, device=self.device)
+        values = T.as_tensor(self.memory.values, dtype=T.float32, device=self.device)
+
+        # Monte Carlo returns (episode-aware)
+        with T.no_grad():
+            returns = T.zeros_like(rewards)
+            G = 0.0
+            for t in reversed(range(rewards.size(0))):
+                G = rewards[t] + self.gamma * G * (1.0 - dones[t])
+                returns[t] = G
+            adv = returns - values
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        dist = self.policy.next_action(states)
+        new_logp = dist.log_prob(actions)
+
+        """PPO Components: Policy update, weighted probability distribution, clipped returns """
+
+        # Updating the policy: update probability distribution (i.e., compute clipped probs)
+        ratio = (new_logp - old_logp).exp()
+
+        # Weighted probaility distribution (according to the formula/update rule)
+        surr1 = ratio * adv
+        surr2 = T.clamp(ratio, 1 - self.clip, 1 + self.clip) * adv
+        value_pred = self.critic.evaluated_state(states)
+
+        # Actor loss: minimize negative of the clipped objective
+        policy_loss = -T.min(surr1, surr2).mean()
+        # Loss: MSE of (return - critic value)
+        value_loss = 0.5 * (returns - value_pred).pow(2).mean()
+        # Entropy (account for randomness in action selection)
+        entropy = dist.entropy().mean()
+        # Total loss: policy loss + constant * value loss - constant * entropy 
+        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()
+
+        self.memory.clear()
+
+
+class Policy(nn.Module):
+    def __init__(self, obs_dim: int, action_dim: int, hidden: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(obs_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, action_dim)
+        )
+
+    def next_action(self, state: T.Tensor) -> Categorical:
+        # Returns the probability distribution over actions
+        if state.dim() == 1:
+            state = state.unsqueeze(0)
+        state = state.view(state.size(0), -1)
+        return Categorical(logits=self.net(state))
+
+
+class Critic(nn.Module): 
+    def __init__(self, obs_dim: int, hidden: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(obs_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, 1)
+        )
+
+    def evaluated_state(self, x: T.Tensor) -> T.Tensor:
+        if x.dim() == 1:
+            x = x.unsqueeze(0)
+        x = x.view(x.size(0), -1)
+        return self.net(x).squeeze(-1)
+
+
+class Memory():
+    def __init__(self):
+        self.states = [] 
+        self.actions = []
+        self.rewards = []
+        self.dones = []
+        self.log_probs = []
+        self.values = []
+        self.next_values = []
+
+    def store(self, state, action, reward, done, log_prob, 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.log_probs.append(float(log_prob))
+        self.values.append(float(value))
+        self.next_values.append(float(next_value))
+
+    """
+    # For mini-batch updates? To be implemented
+    def start_batch(self, batch_size: int):
+        n_states = len(self.states)
+        starts = np.arange(0, n_states, batch_size)
+        index = np.arange(n_states, dtype=np.int64)
+        np.random.shuffle(index)
+        return [index[s:s + batch_size] for s in starts]
+    """
+    def clear(self):
+        self.states = []
+        self.actions = []
+        self.rewards = []
+        self.dones = []
+        self.log_probs = []
+        self.values = []
+        self.next_values = []

Adrian/template_v2_ppo (1).py

@@ -0,0 +1,68 @@
+import ale_py
+import gymnasium as gym
+import sys
+import numpy as np
+from ppo_helpers_v2 import *
+
+def preprocess(obs):
+    # Flatten and normalize uint8 frames to float32 in [0,1]
+    return (obs.astype(np.float32).ravel() / 255.0)
+
+def main() -> int:
+    # Initialize environment
+    env = gym.make("ALE/Pacman-v5", render_mode="human")  # consider removing render_mode for training speed
+    # Initialize variables
+    episode = 0
+    total_return = 0
+    ep_return = 0
+    steps = 2000 # Batch of 100, 2000 environment steps
+    batches = 100
+
+    # Inspect spaces
+    print("Observation space:", env.observation_space)
+    print("Action space:", env.action_space)
+
+    # Create PPO Agent (adapted to ppo_helpers_v2.Agent signature)
+    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, value_coef=0.5, seed=70)
+
+    try:
+        obs, info = env.reset(seed=42)
+        state = preprocess(obs)
+
+        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)
+
+            agent._update()
+            avg_ret = (total_return / episode) if episode else 0
+            print(f"Update {update}: episodes={episode}, avg_return={avg_ret:.2f}")
+
+    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())