Preprocessing visualisation for Enduro and Breakout.

hallvardnmbu · Mar 12, 2024 · c0631ae · c0631ae
1 parent 08302f4
commit c0631ae
Show file tree

Hide file tree

Showing 3 changed files with 251 additions and 2 deletions.
diff --git a/breakout/train.py b/breakout/train.py
@@ -200,7 +200,7 @@
 
     if TRAINING and game % CHECKPOINT == 0:
         logger.info("Saving model")
-        torch.save(value_agent, f"./output/checkpoint-{game}.pth")
+        torch.save(value_agent.state_dict(), f"./output/weights-{game}.pth")
 
 logger.info("Total training time: %s seconds", round(time.time() - start, 2))
 logger.debug("Metrics saved to %s", METRICS)
diff --git a/enduro/train.py b/enduro/train.py
@@ -193,7 +193,7 @@
 
     if TRAINING and game % CHECKPOINT == 0:
         logger.info("Saving model")
-        torch.save(value_agent, f"./output/checkpoint-{game}.pth")
+        torch.save(value_agent.state_dict(), f"./output/weights-{game}.pth")
 
 logger.info("Total training time: %s seconds", round(time.time() - start, 2))
 logger.debug("Metrics saved to %s", METRICS)