From e26564e0ec038c4788c6634596ecaa25a5630696 Mon Sep 17 00:00:00 2001 From: Noah Dormann Date: Thu, 28 Nov 2019 13:35:16 +0100 Subject: [PATCH] Added function for setting up any attributes that weren't saved and thus not loaded --- tests/test_save_load.py | 29 +++++++++++++----------- torchy_baselines/common/base_class.py | 32 ++++++++++++++++++++++++--- torchy_baselines/sac/sac.py | 9 ++++++++ torchy_baselines/td3/td3.py | 11 ++++++++- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/tests/test_save_load.py b/tests/test_save_load.py index e184c6e..4d7861e 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -10,9 +10,9 @@ from torchy_baselines.common.identity_env import IdentityEnvBox MODEL_LIST = [ PPO, - #A2C, - #TD3, - #SAC, + A2C, + TD3, + 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 " \ - "after changing model parameters." + assert not th.allclose(params[k], new_params[k]), "Parameters did not change as expected." 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(): - for step_entry, entry_dict in opt_state['state'].items(): - for value_key,value in entry_dict.items(): - print(value == new_opt_params[optimizer][step_entry][value_key]) - - - + # comparing states not implemented so far. hashes of state_entries are not the same for equal tensors + # comparing every sub_entry does not work because of bool value of Tensor with more than one value is ambiguous + # so far only comparing param_lists + 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) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index fed3243..970ea5a 100644 --- a/torchy_baselines/common/base_class.py +++ b/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) \ No newline at end of file diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index ad36f64..42771bb 100644 --- a/torchy_baselines/sac/sac.py +++ b/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 diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index e21c65a..07036cf 100644 --- a/torchy_baselines/td3/td3.py +++ b/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) \ No newline at end of file