Adding checkpoints: saving the model every n episodes

Author: xeviknal

helpers.py

@@ -1,4 +1,5 @@
 from pyvirtualdisplay import Display
+import torch
 import glob
 import io
 import base64
@@ -23,3 +24,7 @@ def show_video():
 def display_start():
     display = Display(visible=0, size=(1400, 900))
     display.start()
+
+
+def save_model(model, path):
+    torch.save(model.state_dict(), path)

params/.gitignore

@@ -0,0 +1 @@
+

trainer.py

@@ -2,6 +2,7 @@
 import numpy as np
 from policy import Policy
 from actions import Actions
+from helpers import save_model
 
 
 class Trainer:
@@ -80,6 +81,8 @@ def train(self):
             if i_episode % self.config['log_interval'] == 0:
                 print('Episode {}\tLast reward: {:.2f}\tAverage reward: {:.2f}'.format(
                     i_episode, ep_reward, running_reward))
+                save_model(self.policy, './params/policy-params.dl')
+
             if running_reward > self.env.spec().reward_threshold:
                 print("Solved!")
                 break