feat: Agent can now save model

urnai/agents/agent_base.py

@@ -57,3 +57,9 @@ def learn(self, obs, reward, done) -> None:
             next_state = self.state_space.update(obs)
             self.model.learn(self.previous_state, self.previous_action,
                               reward, next_state, done)
+
+    def save(self, savepath) -> None:
+        if (self.model.persistence == None):
+            raise 'No persistence set in model'
+
+        self.model.persistence.save(savepath)

urnai/base/persistence_pickle.py

@@ -108,7 +108,7 @@ def _get_attributes(self):
     def _get_dict(self):
         pickleable_attr_dict = {}
 
-        for attr in self.object_to_save._get_attributes():
+        for attr in self._get_attributes():
             pickleable_attr_dict[attr] = getattr(self.object_to_save, attr)
 
         return pickleable_attr_dict

urnai/models/model_base.py

@@ -8,6 +8,7 @@ class ModelBase(ABC):
 
 	def __init__(self):
 		self.learning_data = {}
+		self.persistence = None
 
 	@abstractmethod
 	def learn(self, current_state, action, reward, next_state):

urnai/trainers/trainer.py

@@ -88,6 +88,7 @@ def training_loop(self, is_training, reward_from_agent=True):
 
                 if done:
                     print("Episode: %d, Reward: %d" % (current_episodes, ep_reward))
+                    self.agent.save("saves/")
                     break
 
             # if this is not a test (evaluation), saving is enabled and we are in a multiple