Added function for setting up any attributes that weren't saved and thus not loaded

tests/test_save_load.py

@@ -10,9 +10,9 @@ from torchy_baselines.common.identity_env import IdentityEnvBox
 
 MODEL_LIST = [
     PPO,
-    #A2C,
+    A2C,
-    #TD3,
+    TD3,
-    #SAC,
+    SAC,
 ]
 
 
@@ -30,7 +30,7 @@ def test_save_load(model_class):
 
     # create model
     model = model_class('MlpPolicy', env, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
-    model.learn(total_timesteps=1000, eval_freq=500)
+    model.learn(total_timesteps=500, eval_freq=250)
 
     # Get dictionary of current parameters
     params = deepcopy(model.get_policy_parameters())
@@ -45,8 +45,7 @@ def test_save_load(model_class):
     new_params = model.get_policy_parameters()
     # Check that all params are different now
     for k in params:
-        assert not th.allclose(params[k], new_params[k]), "Selected actions did not change " \
+        assert not th.allclose(params[k], new_params[k]), "Parameters did not change as expected."
-                                                          "after changing model parameters."
 
     params = new_params
 
@@ -67,13 +66,17 @@ def test_save_load(model_class):
     # check if keys are the same
     assert opt_params.keys() == new_opt_params.keys()
     # check if values are the same: only tested for Adam and RMSProp so far
-    for optimizer,opt_state in opt_params.items():
+    # comparing states not implemented so far. hashes of state_entries are not the same for equal tensors
-        for step_entry, entry_dict in opt_state['state'].items():
+    # comparing every sub_entry does not work because of bool value of Tensor with more than one value is ambiguous
-            for value_key,value in entry_dict.items():
+    # so far only comparing param_lists
-                print(value == new_opt_params[optimizer][step_entry][value_key])
+    for optimizer, opt_state in opt_params.items():
-
+        for param_group_idx, param_group in enumerate(opt_state['param_groups']):
-
+            for param_key, param_value in param_group.items():
-
+                if param_key == 'params':  # don't know how to handle params correctly, therefore only check if we have the same amount
+                    assert len(param_value) == len(
+                        new_opt_params[optimizer]['param_groups'][param_group_idx][param_key])
+                else:
+                    assert param_value == new_opt_params[optimizer]['param_groups'][param_group_idx][param_key]
 
     # check if learn still works
     model.learn(total_timesteps=1000, eval_freq=500)

torchy_baselines/common/base_class.py

@@ -292,7 +292,8 @@ class BaseRLModel(object):
         model.__dict__.update(kwargs)
         model.set_env(env)
         model.load_parameters(params, opt_params)
-
+        # resetup modul after load
+        model._resetup_model()
         return model
 
     @staticmethod
@@ -511,15 +512,40 @@ class BaseRLModel(object):
                     with file_.open(file_name + '.pth', mode="w") as opt_param_file:
                         th.save(dict, opt_param_file)
 
-    def save(self, path, include=None):#TODO
+    def excluded_save_params(self):
+        """
+        returns the names of the parameters that should be excluded from save
+        :return: (list) List of parameters that should be excluded from save
+        """
+        return ["replay_buffer"]
+
+    def _resetup_model(self):
+        """
+        Function will be called at the end of load and should resetup anything that might not have been saved
+        warning: this function should always be in compliance with excluded_save_params
+        :return: 
+        """
+        pass
+
+    def save(self, path, include=None):
         """
         saves all the params from init and pytorch params in a file for continuous learning
 
         :param path: (str) path to the file where the data should be saved
+        :param include: (list) name of parameters that might be excluded but should be included anyway
         :return:
         """
         data = self.__dict__
-        data.pop("replay_buffer")
+        # get list of params to be excluded
+        exclude = self.excluded_save_params()
+        # do not exclude params if they are specifically included
+        if include is not None:
+            exclude = [param_name for param_name in exclude if param_name not in include]
+
+        # remove parameter entries of parameters which are to be excluded
+        for param_name in exclude:
+            data.pop(param_name, None)
+
         params_to_save = self.get_policy_parameters()
         opt_params_to_save = self.get_opt_parameters()
         self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)

torchy_baselines/sac/sac.py

@@ -128,6 +128,15 @@ class SAC(BaseRLModel):
         self.policy = self.policy.to(self.device)
         self._create_aliases()
 
+    def _resetup_model(self):
+        """
+        method used to resetup anything that was not saved 
+        :return: 
+        """
+        if self.replay_buffer is None:
+            obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
+            self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
+
     def _create_aliases(self):
         self.actor = self.policy.actor
         self.critic = self.policy.critic

torchy_baselines/td3/td3.py

@@ -85,6 +85,15 @@ class TD3(BaseRLModel):
         self.policy = self.policy.to(self.device)
         self._create_aliases()
 
+    def _resetup_model(self):
+        """
+        method used to resetup anything that was not saved 
+        :return: 
+        """
+        if self.replay_buffer is None:
+            obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
+            self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
+
     def _create_aliases(self):
         self.actor = self.policy.actor
         self.actor_target = self.policy.actor_target
@@ -256,4 +265,4 @@ class TD3(BaseRLModel):
         """
         self.actor.optimizer.load_state_dict(opt_params["actor"])
         self.critic.optimizer.load_state_dict(opt_params["critic"])
-        self.policy.load_state_dict(load_dict)
+        self.policy.load_state_dict(load_dict)