Update training.py

training.py

@@ -251,41 +251,58 @@ def supervised_warmup(model, tokenizer, data_path="training_data.json", epochs=3
 # =========================================================
 # PPO UPDATE (FIXED advantage = return – baseline)
 # =========================================================
-def ppo_update(trajectories, model, tokenizer, optimizer, clip=0.2):
+def ppo_update(trajectories, model, tokenizer, optimizer, clip=0.2, gamma=0.99):
     model.train()
 
     losses = []
     kls = []
 
-    # Gather all returns and compute a global baseline (simple REINFORCE)
+    # =========================
+    # Compute returns + baseline
+    # =========================
     all_returns = []
-    for traj in trajectories:
-        returns = np.cumsum(traj.rewards[::-1])[::-1]
-        all_returns.extend(returns)
-    baseline = np.mean(all_returns) if all_returns else 0.0
 
+    traj_returns = []
     for traj in trajectories:
-        returns = np.cumsum(traj.rewards[::-1])[::-1]
-        returns = torch.tensor(returns, device=DEVICE)
+        returns = []
+        running = 0.0
+
+        for r in reversed(traj.rewards):
+            running = r + gamma * running
+            returns.insert(0, running)
+
+        returns = torch.tensor(returns, dtype=torch.float32, device=DEVICE)
+        traj_returns.append(returns)
+        all_returns.extend(returns.tolist())
+
+    baseline = torch.tensor(np.mean(all_returns), device=DEVICE) if all_returns else torch.tensor(0.0, device=DEVICE)
+
+    # =========================
+    # PPO update
+    # =========================
+    for traj, returns in zip(trajectories, traj_returns):
 
         for i in range(len(traj.states)):
             state = traj.states[i]
             action = traj.actions[i]
+
             old_lp = torch.tensor(traj.logprobs[i], device=DEVICE)
 
-            # Proper advantage: return – baseline
-            adv = returns[i] - baseline
+            # Advantage (detached)
+            adv = (returns[i] - baseline).detach()
 
             messages = [{"role": "user", "content": state}]
-            formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-            full = formatted + action
+            formatted = tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=True
+            )
 
+            full = formatted + action
             inputs = tokenizer(full, return_tensors="pt", truncation=True).to(DEVICE)
+
             logits = model(**inputs).logits
 
             action_ids = tokenizer.encode(action, add_special_tokens=False)
-            prefix_ids = tokenizer.encode(formatted, add_special_tokens=False)
-            prefix_len = len(prefix_ids)
+            prefix_len = len(tokenizer.encode(formatted, add_special_tokens=False))
 
             logps = []
             entropy = 0.0
@@ -295,8 +312,9 @@ def ppo_update(trajectories, model, tokenizer, optimizer, clip=0.2):
                 if pos == 0 or pos >= logits.shape[1]:
                     continue
 
-                token_logits = logits[0, pos-1]
+                token_logits = logits[0, pos - 1]
                 log_probs = F.log_softmax(token_logits, dim=-1)
+
                 lp = log_probs[action_ids[idx]]
                 logps.append(lp)
 
@@ -307,10 +325,15 @@ def ppo_update(trajectories, model, tokenizer, optimizer, clip=0.2):
                 continue
 
             new_lp = torch.stack(logps).sum()
+
+            # PPO ratio
             ratio = torch.exp(new_lp - old_lp)
+
             s1 = ratio * adv
-            s2 = torch.clamp(ratio, 1-clip, 1+clip) * adv
-            loss = -torch.min(s1, s2) - 0.01 * entropy
+            s2 = torch.clamp(ratio, 1 - clip, 1 + clip) * adv
+
+            policy_loss = -torch.min(s1, s2)
+            loss = policy_loss - 0.01 * (entropy / len(logps))
 
             if torch.isnan(loss):
                 continue
@@ -324,8 +347,10 @@ def ppo_update(trajectories, model, tokenizer, optimizer, clip=0.2):
             kls.append(kl)
             losses.append(loss.item())
 
-    return np.mean(losses) if losses else 0.0, np.mean(kls) if kls else 0.0
-
+    return (
+        float(np.mean(losses)) if losses else 0.0,
+        float(np.mean(kls)) if kls else 0.0,
+    )
 # =========================================================
 # MAIN TRAINING LOOP
 # =========================================================