Independent save/load for policies

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:
 ^^^^^^^^^^

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")

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,

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)