From f347474e6a778fa5a944cff2d9cca73640ccc858 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 20 Apr 2020 15:59:44 +0200 Subject: [PATCH] Independent save/load for policies --- docs/misc/changelog.rst | 5 +++- tests/test_save_load.py | 25 ++++++++++------- torchy_baselines/common/policies.py | 2 +- torchy_baselines/sac/policies.py | 43 +++++++++++++++++++++++++++++ 4 files changed, 63 insertions(+), 12 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 13c7ca5..41bf91e 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,16 +3,19 @@ Changelog ========== -Pre-Release 0.5.0a0 (WIP) +Pre-Release 0.5.0a1 (WIP) ------------------------------ Breaking Changes: ^^^^^^^^^^^^^^^^^ +- Previous loading of policy weights is broken and replace by the new saving/loading for policy New Features: ^^^^^^^^^^^^^ - Added ``optimizer`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily customizer optimizers +- Complete independent save/load for policies + Bug Fixes: ^^^^^^^^^^ diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 3bc7d23..1e547c0 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -180,9 +180,11 @@ def test_save_load_policy(model_class): observations = observations.reshape(10, -1) policy = model.policy - actor = None + policy_class = policy.__class__ + actor, actor_class = None, None if model_class in [SAC, TD3]: actor = policy.actor + actor_class = actor.__class__ # Get dictionary of current parameters params = deepcopy(policy.state_dict()) @@ -207,14 +209,17 @@ def test_save_load_policy(model_class): selected_actions_actor, _ = actor.predict(observations, deterministic=True) # Save and load policy - policy.save("./logs/policy_weights.pkl") + policy.save("./logs/policy.pkl") # Save and load actor if actor is not None: - actor.save("./logs/actor_weights.pkl") + actor.save("./logs/actor.pkl") - policy.load("./logs/policy_weights.pkl") - if actor is not None: - actor.load("./logs/actor_weights.pkl") + del policy, actor + + + policy = policy_class.load("./logs/policy.pkl") + if actor_class is not None: + actor = actor_class.load("./logs/actor.pkl") # check if params are still the same after load new_params = policy.state_dict() @@ -227,12 +232,12 @@ def test_save_load_policy(model_class): new_selected_actions, _ = policy.predict(observations, deterministic=True) assert np.allclose(selected_actions, new_selected_actions, 1e-4) - if actor is not None: + if actor_class is not None: new_selected_actions_actor, _ = actor.predict(observations, deterministic=True) assert np.allclose(selected_actions_actor, new_selected_actions_actor, 1e-4) assert np.allclose(selected_actions_actor, new_selected_actions, 1e-4) # clear file from os - os.remove("./logs/policy_weights.pkl") - if actor is not None: - os.remove("./logs/actor_weights.pkl") + os.remove("./logs/policy.pkl") + if actor_class is not None: + os.remove("./logs/actor.pkl") diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 982a400..41b7211 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -203,7 +203,7 @@ class BasePolicy(nn.Module): return dict( observation_space=self.observation_space, action_space=self.action_space, - # Passed to the constructor by child classes + # Passed to the constructor by child class # squash_output=self.squash_output, # features_extractor=self.features_extractor normalize_images=self.normalize_images, diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index dd128bb..c0fd3a6 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -66,6 +66,15 @@ class Actor(BasePolicy): self.latent_pi = nn.Sequential(*latent_pi_net) self.use_sde = use_sde self.sde_features_extractor = None + self.sde_net_arch = sde_net_arch + self.net_arch = net_arch + self.features_dim = features_dim + self.activation_fn = activation_fn + self.log_std_init = log_std_init + self.sde_net_arch = sde_net_arch + self.use_expln = use_expln + self.full_std = full_std + self.clip_mean = clip_mean if self.use_sde: latent_sde_dim = net_arch[-1] @@ -88,6 +97,23 @@ class Actor(BasePolicy): self.mu = nn.Linear(net_arch[-1], action_dim) self.log_std = nn.Linear(net_arch[-1], action_dim) + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_arch, + features_dim=self.features_dim, + activation_fn=self.activation_fn, + use_sde=self.use_sde, + log_std_init=self.log_std_init, + full_std=self.full_std, + sde_net_arch=self.sde_net_arch, + use_expln=self.use_expln, + features_extractor=self.features_extractor, + clip_mean=self.clip_mean + )) + return data + def get_std(self) -> th.Tensor: """ Retrieve the standard deviation of the action distribution. @@ -286,6 +312,23 @@ class SACPolicy(BasePolicy): self.critic.optimizer = self.optimizer_class(self.critic.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs) + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update(dict( + net_arch=self.net_args['net_arch'], + activation_fn=self.net_args['activation_fn'], + use_sde=self.actor_kwargs['use_sde'], + log_std_init=self.actor_kwargs['log_std_init'], + sde_net_arch=self.actor_kwargs['sde_net_arch'], + use_expln=self.actor_kwargs['use_expln'], + clip_mean=self.actor_kwargs['clip_mean'], + lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone + optimizer=self.optimizer_class, + optimizer_kwargs=self.optimizer_kwargs + )) + return data + def make_actor(self) -> Actor: return Actor(**self.actor_kwargs).to(self.device)