From 517c17047b7d9e7cbc510464754f6178b0b2c73d Mon Sep 17 00:00:00 2001
From: rgriffith <KoalaTrapLord@users.noreply.github.com>
Date: Mon, 30 Mar 2026 13:31:59 -0400
Subject: [PATCH] fix: strip value_head keys when loading PPO checkpoint into
 PolicyNetwork

---
 ml/training/policy_network.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/ml/training/policy_network.py b/ml/training/policy_network.py
index b84a7dd..8500b3e 100644
--- a/ml/training/policy_network.py
+++ b/ml/training/policy_network.py
@@ -223,7 +223,8 @@ def load_checkpoint(path: str, device: str = "cpu"):
         num_layers=config.get("num_layers", TRAINING_CONFIG.num_layers),
         dropout=config.get("dropout", TRAINING_CONFIG.dropout),
     )
-    model.load_state_dict(checkpoint["model_state_dict"])
+    state = {k: v for k, v in checkpoint["model_state_dict"].items() if not k.startswith("value_head")}
+    model.load_state_dict(state, strict=False)
     model.to(device)
 
     return model, checkpoint