From 4b6234a1c84384aa3a5c3be377006b1a4f749cb2 Mon Sep 17 00:00:00 2001
From: Noah Dormann <Noah.Dormann@dlr.de>
Date: Thu, 21 Nov 2019 11:39:47 +0100
Subject: [PATCH] finished test_save_load.py test

---
 tests/test_save_load.py               | 31 +++++++++++++++++----------
 torchy_baselines/common/base_class.py | 13 ++++++++---
 torchy_baselines/ppo/ppo.py           |  2 +-
 3 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/tests/test_save_load.py b/tests/test_save_load.py
index ca3f637..9941e15 100644
--- a/tests/test_save_load.py
+++ b/tests/test_save_load.py
@@ -1,8 +1,9 @@
 import os
 
 import pytest
+import copy
 import numpy as np
-
+import torch
 from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
 from torchy_baselines.common.noise import NormalActionNoise
 from torchy_baselines.common.vec_env import DummyVecEnv
@@ -17,6 +18,7 @@ MODEL_LIST = [
 def test_save_load(model_class):
     """
     Test if 'save' and 'load' saves and loads model correctly
+    and if 'load_parameters' and 'get_policy_parameters' work correctly
 
     :param model_class: (BaseRLModel) A RL model
     """
@@ -26,25 +28,32 @@ def test_save_load(model_class):
     model = model_class('MlpPolicy', env, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
 
     # test action probability for given (obs, action) pair
-    env = model.get_env()
-    obs = env.reset()
-    observations = np.array([obs for _ in range(10)])
-    observations = np.squeeze(observations)
-
-    #actions = np.array([env.action_space.sample() for _ in range(10)])
 
     # Get dictionary of current parameters
-    params = model.get_parameters()
+    params = copy.deepcopy(model.get_policy_parameters())
 
     # Modify all parameters to be random values
-    random_params = dict((param_name,np.random.random(size=param.shape)) for param_name, param in params.items())
+    random_params = dict((param_name, torch.rand_like(param)) for param_name, param in params.items())
     # Update model parameters with the new zeroed values
     model.load_parameters(random_params)
-    # Get new action probas
-    #...
+
+    # shared items
+    new_params = model.get_policy_parameters()
+    shared_items = {k: params[k] for k in params if k in new_params and torch.all(torch.eq(params[k], new_params[k]))}
+    # Check that at least some actions are chosen different now
+    assert not len(shared_items) == len(new_params), "Selected actions did not change " \
+                                                     "after changing model parameters."
+
+    params = new_params
 
     # Check
     model.learn(total_timesteps=1000, eval_freq=500)
     model.save("test_save.zip")
     model = model.load("test_save")
+
+    #check if params are still the same after load
+    new_params = model.get_policy_parameters()
+    shared_items = {k: params[k] for k in params if k in new_params and torch.all(torch.eq(params[k], new_params[k]))}
+    # Check that at least some actions are chosen different now
+    assert len(shared_items) == len(new_params), "Parameters not the same after save and load."
     os.remove("test_save.zip")
diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py
index fd85b0a..d0c0d5d 100644
--- a/torchy_baselines/common/base_class.py
+++ b/torchy_baselines/common/base_class.py
@@ -186,14 +186,21 @@ class BaseRLModel(object):
         """
         return self.params
 
-    def get_parameters(self):
+    def get_policy_parameters(self):
         """
-        Get current model parameters as dictionary of variable name -> ndarray.
+        Get current model policy parameters as dictionary of variable name -> tensors.
 
-        :return: (OrderedDict) Dictionary of variable name -> ndarray of model's parameters.
+        :return: (OrderedDict) Dictionary of variable name -> tensor of model's policy parameters.
         """
         return self.policy.state_dict()
 
+    def get_optim_parameters(self):
+        """
+        Get current model optimizer parameters as dictionary of variable names -> tensors
+        :return: (OrderedDict) Dictionary of variable name -> tensor of model's optimizer parameters
+        """
+        raise NotImplementedError()
+
     def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
                  adam_epsilon=1e-8, val_interval=None):
         """
diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py
index e6467f3..f22297b 100644
--- a/torchy_baselines/ppo/ppo.py
+++ b/torchy_baselines/ppo/ppo.py
@@ -319,7 +319,7 @@ class PPO(BaseRLModel):
 
         }
 
-        params_to_save = self.get_parameters()
+        params_to_save = self.get_policy_parameters()
 
         self._save_to_file_zip(path, data=data, params=params_to_save)