Added vanilla_ppo_update (base case w/o fancy normalizations)

CNN_PPO/ppo_helpers_cnn.py

@@ -106,6 +106,89 @@ class Agent:
             next_value = self.critic.evaluated_state(ns).item()
         self.memory.store(state, action, reward, done, log_prob, value, next_value)
 
+    def vanilla_ppo_update(self):
+        if len(self.memory.states) == 0:
+            return 0.0
+
+        # 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)
+        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)
+
+        with T.no_grad():
+            # Compute next values (bootstrap for final step)
+            next_values = T.cat([values[1:], values[-1:].clone()])
+            deltas = rewards + self.gamma * next_values * (1 - dones) - values
+
+            # --- GAE-Lambda ---
+            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
+            # Advantage normalization
+            adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
+
+        # --- PPO Multiple Epochs + Minibatch ---
+        total_loss_epoch = 0.0
+        num_samples = len(states)
+        batch_size = min(64, num_samples)
+        ppo_epochs = 4
+
+        for _ in range(ppo_epochs):
+            # Shuffle indices
+            idxs = T.randperm(num_samples)
+            for start in range(0, num_samples, batch_size):
+                batch_idx = idxs[start:start + batch_size]
+
+                b_states = states[batch_idx]
+                b_actions = actions[batch_idx]
+                b_old_logp = old_logp[batch_idx]
+                b_returns = returns[batch_idx]
+                b_adv = adv[batch_idx]
+
+                dist = self.policy.next_action(b_states)
+                new_logp = dist.log_prob(b_actions)
+                entropy = dist.entropy().mean()
+                ratio = (new_logp - b_old_logp).exp()
+
+                # --- Clipped surrogate objective ---
+                surr1 = ratio * b_adv
+                surr2 = T.clamp(ratio, 1 - self.clip, 1 + self.clip) * b_adv
+                policy_loss = -T.min(surr1, surr2).mean()
+
+                # --- Critic loss ---
+                value_pred = self.critic.evaluated_state(b_states)
+                value_loss = 0.5 * (b_returns - value_pred).pow(2).mean()
+
+                # --- Total loss ---
+                total_loss = (
+                        policy_loss +
+                        self.value_coef * value_loss -
+                        self.entropy_coef * entropy
+                )
+
+                # Debug: track individual loss components
+                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()
+                self.opt.step()
+
+                total_loss_epoch += total_loss.item()
+            
+        # Clear memory after full PPO update
+        self.memory.clear()
+
+        return total_loss_epoch / (ppo_epochs * (num_samples / batch_size))
+
+
     def update_rbs(self):
         if len(self.memory.states) == 0:
             return 0.0
@@ -135,9 +218,10 @@ class Agent:
             # --- Return-based normalization (RBS) ---
             sigma_t = returns.std(unbiased=False) + 1e-8
             returns = returns / sigma_t
+            self.sigma_history.append(sigma_t.item())
             adv = adv / sigma_t
+            # Advantage normalization
             adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
-            self.sigma_history.append(sigma_t.item())
 
         # --- PPO Multiple Epochs + Minibatch ---
         total_loss_epoch = 0.0
@@ -178,6 +262,10 @@ class Agent:
                         self.entropy_coef * entropy
                 )
 
+                # Debug: track individual loss components
+                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()
                 self.opt.step()
@@ -216,13 +304,8 @@ class Agent:
                 adv[t] = gae
 
             returns = adv + values
-
-            # --- Return-based normalization (RBS) ---
-            sigma_t = returns.std(unbiased=False) + 1e-8
-            returns = returns / sigma_t
-            adv = adv / sigma_t
+            # Advantage normalization
             adv = (adv - adv.mean()) / (adv.std(unbiased=False) + 1e-8)
-            self.sigma_history.append(sigma_t.item())
 
         # --- PPO Multiple Epochs + Minibatch ---
         total_loss_epoch = 0.0
@@ -262,6 +345,10 @@ class Agent:
                         self.value_coef * value_loss -
                         self.entropy_coef * entropy
                 )
+                
+                # Debug: track individual loss components
+                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()
@@ -353,11 +440,9 @@ class Critic(nn.Module):
         c, h, w = obs_shape
         # Suggested architecture for Atari: https://arxiv.org/pdf/1312.5602
         self.cnn = nn.Sequential(
-            nn.Conv2d(c, 32, kernel_size=8, stride=4),
-            nn.ReLU(),
-            nn.Conv2d(32, 64, kernel_size=4, stride=2),
+            nn.Conv2d(c, 16, kernel_size=8, stride=4),
             nn.ReLU(),
-            nn.Conv2d(64, 64, kernel_size=3, stride=1),
+            nn.Conv2d(16, 32, kernel_size=4, stride=2),
             nn.ReLU(),
             nn.Flatten()
         )