first save and load features

tests/test_run.py

@@ -19,6 +19,16 @@ def test_td3():
     os.remove("test_save.pth")
 
 
+
+
+def test_a2c():
+    model = A2C('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
+    model.learn(total_timesteps=1000, eval_freq=500)
+    model.save("test_save")
+    model.load("test_save")
+    os.remove("test_save.pth")
+
+
 def test_cemrl():
     model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1,
                   learning_starts=100, verbose=1, create_eval_env=True, action_noise=action_noise)

tests/test_save_load.py

@@ -0,0 +1,50 @@
+import os
+
+import pytest
+import numpy as np
+
+from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
+from torchy_baselines.common.noise import NormalActionNoise
+from torchy_baselines.common.vec_env import DummyVecEnv
+from torchy_baselines.common.identity_env import IdentityEnvBox
+
+MODEL_LIST = [
+    PPO
+]
+
+
+@pytest.mark.parametrize("model_class", MODEL_LIST)
+def test_save_load(model_class):
+    """
+    Test if 'save' and 'load' saves and loads model correctly
+
+    :param model_class: (BaseRLModel) A RL model
+    """
+    env = DummyVecEnv([lambda: IdentityEnvBox(10)])
+
+    # create model
+    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()
+
+    # 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())
+    # Update model parameters with the new zeroed values
+    model.load_parameters(random_params)
+    # Get new action probas
+    #...
+
+    # Check
+    model.learn(total_timesteps=1000, eval_freq=500)
+    model.save("test_save.zip")
+    model = model.load("test_save")
+    os.remove("test_save.zip")

torchy_baselines/a2c/a2c.py

@@ -124,3 +124,15 @@ class A2C(PPO):
         return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval,
                   eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
                   tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
+
+    def save(self, path):
+        if not path.endswith('.pth'):
+            path += '.pth'
+        th.save(self.policy.state_dict(), path)
+
+    def load(self, path, env=None, **_kwargs):
+        if not path.endswith('.pth'):
+            path += '.pth'
+        if env is not None:
+            pass
+        self.policy.load_state_dict(th.load(path))

torchy_baselines/common/base_class.py

@@ -12,6 +12,12 @@ from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv
 from torchy_baselines.common.monitor import Monitor
 from torchy_baselines.common import logger
 
+# for storing and loging
+import os
+import io
+import zipfile
+from torchy_baselines.common.save_util import (data_to_json, json_to_data)
+
 
 class BaseRLModel(object):
     """
@@ -57,6 +63,7 @@ class BaseRLModel(object):
         self.replay_buffer = None
         self.seed = seed
         self.action_noise = None
+        self.params = None
         # Track the training progress (from 1 to 0)
         # this is used to update the learning rate
         self._current_progress = 1
@@ -113,7 +120,7 @@ class BaseRLModel(object):
         (no need for symmetric action space)
         """
         low, high = self.action_space.low, self.action_space.high
-        return low + (0.5 * (scaled_action + 1.0) * (high -  low))
+        return low + (0.5 * (scaled_action + 1.0) * (high - low))
 
     def _setup_learning_rate(self):
         """Transform to callable if needed."""
@@ -179,7 +186,7 @@ class BaseRLModel(object):
 
         :return: (list) List of pytorch Variables
         """
-        pass
+        return self.params
 
     def get_parameters(self):
         """
@@ -187,7 +194,7 @@ class BaseRLModel(object):
 
         :return: (OrderedDict) Dictionary of variable name -> ndarray of model's parameters.
         """
-        raise NotImplementedError()
+        return self.policy.state_dict()
 
     def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
                  adam_epsilon=1e-8, val_interval=None):
@@ -237,14 +244,11 @@ class BaseRLModel(object):
         """
         pass
 
-    def load_parameters(self, load_path_or_dict, exact_match=True):
+    def load_parameters(self, load_dict, exact_match=True):
         """
-        Load model parameters from a file or a dictionary
+        Load model parameters from a dictionary
 
-        Dictionary keys should be tensorflow variable names, which can be obtained
-        with ``get_parameters`` function. If ``exact_match`` is True, dictionary
-        should contain keys for all model's parameters, otherwise RunTimeError
-        is raised. If False, only variables included in the dictionary will be updated.
+        Dictionary should be of shape torch model.state_dict()
 
         This does not load agent's hyper-parameters.
 
@@ -252,13 +256,10 @@ class BaseRLModel(object):
             This function does not update trainer/optimizer variables (e.g. momentum).
             As such training after using this function may lead to less-than-optimal results.
 
-        :param load_path_or_dict: (str or file-like or dict) Save parameter location
-            or dict of parameters as variable.name -> ndarrays to be loaded.
-        :param exact_match: (bool) If True, expects load dictionary to contain keys for
-            all variables in the model. If False, loads parameters only for variables
-            mentioned in the dictionary. Defaults to True.
+
+        :param load_path_or_dict: (dict) dict of parameters from model.state_dict()
         """
-        raise NotImplementedError()
+        self.policy.load_state_dict(load_dict)
 
     @abstractmethod
     def save(self, save_path):
@@ -280,7 +281,103 @@ class BaseRLModel(object):
             (can be None if you only need prediction from a trained model)
         :param kwargs: extra arguments to change the model when loading
         """
-        raise NotImplementedError()
+        data, params = cls._load_from_file(load_path)
+
+        if 'policy_kwargs' in kwargs and kwargs['policy_kwargs'] != data['policy_kwargs']:
+            raise ValueError("The specified policy kwargs do not equal the stored policy kwargs."
+                             "Stored kwargs: {}, specified kwargs: {}".format(data['policy_kwargs'],kwargs['policy_kwargs']))
+
+        model = cls(policy=data["policy"],env=None, _init_setup_model=False)
+        model.__dict__.update(data)
+        model.__dict__.update(kwargs)
+        model.set_env(env)
+        model.load_parameters(params)
+
+        return model
+
+
+    @staticmethod
+    def _save_to_file_zip(save_path, data=None, params=None):
+        """Save model to a zip archive
+
+        :param save_path: (str or file-like) Where to store the model
+        :param data: (OrderedDict) Class parameters being stored
+        :param params: (OrderedDict) Model parameters being stored expexted to be state_dict
+        """
+
+        # data/params can be None, so do not
+        # try to serialize them blindly
+        if data is not None:
+            serialized_data = data_to_json(data)
+
+        # Check postfix if save_path is a string
+        if isinstance(save_path, str):
+            _, ext = os.path.splitext(save_path)
+            if ext == "":
+                save_path += ".zip"
+
+        # Create a zip-archive and write our objects
+        # there. This works when save_path is either
+        # str or a file-like
+        with zipfile.ZipFile(save_path, "w") as file_:
+            # Do not try to save "None" elements
+            if data is not None:
+                file_.writestr("data",serialized_data)
+            if params is not None:
+                with file_.open('param.pth', mode="w") as param_file:
+                    th.save(params,param_file)
+
+    @staticmethod
+    def _load_from_file(load_path, load_data = True):
+        """ Load model data from a .zip archive
+
+        :param load_path: (str or file-like) Where to load the model from
+        :param load_data: (bool) Whether we should load and return data
+            (class parameters). Mainly used by 'load_parameters' to only load model parameters (weights)
+        :return: (dict. OrderedDict),(dict. OrderedDict) Class parameters and model parameters (state_dict)
+        """
+        # Check if file exists if load_path is a string
+        if isinstance(load_path, str):
+            if not os.path.exists(load_path):
+                if os.path.exists(load_path + ".zip"):
+                    load_path += ".zip"
+                else:
+                    raise ValueError("Error: the file {} could not be found".format(load_path))
+
+        # Open the zip archive and load data
+        try:
+            with zipfile.ZipFile(load_path,"r") as file_:
+                namelist = file_.namelist()
+                # If data or parameters is not in the
+                # zip archive, assume they were stored
+                # as None (_save_to_file_zip allows this).
+                data = None
+                params = None
+                if "data" in namelist and load_data:
+                    # Load class parameters and convert to string
+                    json_data = file_.read("data").decode()
+                    data = json_to_data(json_data)
+
+                if "param.pth" in namelist:
+                    # Load parameters with build in torch function
+                    with file_.open("param.pth", mode="r") as param_file:
+                        # File has to be seekable so load in BytesIO first
+                        file_content = io.BytesIO()
+                        file_content.write(param_file.read())
+                        # go to start of file
+                        file_content.seek(0)
+                        params = th.load(file_content)
+        except zipfile.BadZipFile:
+            # load_path wasn't a zip file
+            raise ValueError("Error: the file {} wasn't a zip-file".format(load_path))
+
+        return data, params
+
+
+
+
+
+
 
     def set_random_seed(self, seed=0):
         set_random_seed(seed, using_cuda=self.device == th.device('cuda'))
@@ -375,7 +472,8 @@ class BaseRLModel(object):
                     action_noise.reset()
 
                 # Display training infos
-                if self.verbose >= 1 and log_interval is not None and (episode_num + total_episodes) % log_interval == 0:
+                if self.verbose >= 1 and log_interval is not None and (
+                        episode_num + total_episodes) % log_interval == 0:
                     fps = int(num_timesteps / (time.time() - self.start_time))
                     logger.logkv("episodes", episode_num + total_episodes)
                     # logger.logkv("mean 100 episode reward", mean_reward)

torchy_baselines/common/identity_env.py

@@ -0,0 +1,105 @@
+import numpy as np
+
+from gym import Env
+from gym.spaces import Discrete, MultiDiscrete, MultiBinary, Box
+
+
+class IdentityEnv(Env):
+    def __init__(self, dim, ep_length=100):
+        """
+        Identity environment for testing purposes
+
+        :param dim: (int) the size of the dimensions you want to learn
+        :param ep_length: (int) the length of each episodes in timesteps
+        """
+        self.action_space = Discrete(dim)
+        self.observation_space = self.action_space
+        self.ep_length = ep_length
+        self.current_step = 0
+        self.dim = dim
+        self.reset()
+
+    def reset(self):
+        self.current_step = 0
+        self._choose_next_state()
+        return self.state
+
+    def step(self, action):
+        reward = self._get_reward(action)
+        self._choose_next_state()
+        self.current_step += 1
+        done = self.current_step >= self.ep_length
+        return self.state, reward, done, {}
+
+    def _choose_next_state(self):
+        self.state = self.action_space.sample()
+
+    def _get_reward(self, action):
+        return 1 if np.all(self.state == action) else 0
+
+    def render(self, mode='human'):
+        pass
+
+
+class IdentityEnvBox(IdentityEnv):
+    def __init__(self, low=-1, high=1, eps=0.05, ep_length=100):
+        """
+        Identity environment for testing purposes
+
+        :param dim: (int) the size of the dimensions you want to learn
+        :param low: (float) the lower bound of the box dim
+        :param high: (float) the upper bound of the box dim
+        :param eps: (float) the epsilon bound for correct value
+        :param ep_length: (int) the length of each episodes in timesteps
+        """
+        super(IdentityEnvBox, self).__init__(1, ep_length)
+        self.action_space = Box(low=low, high=high, shape=(1,), dtype=np.float32)
+        self.observation_space = self.action_space
+        self.eps = eps
+        self.reset()
+
+    def reset(self):
+        self.current_step = 0
+        self._choose_next_state()
+        return self.state
+
+    def step(self, action):
+        reward = self._get_reward(action)
+        self._choose_next_state()
+        self.current_step += 1
+        done = self.current_step >= self.ep_length
+        return self.state, reward, done, {}
+
+    def _choose_next_state(self):
+        self.state = self.observation_space.sample()
+
+    def _get_reward(self, action):
+        return 1 if (self.state - self.eps) <= action <= (self.state + self.eps) else 0
+
+
+class IdentityEnvMultiDiscrete(IdentityEnv):
+    def __init__(self, dim, ep_length=100):
+        """
+        Identity environment for testing purposes
+
+        :param dim: (int) the size of the dimensions you want to learn
+        :param ep_length: (int) the length of each episodes in timesteps
+        """
+        super(IdentityEnvMultiDiscrete, self).__init__(dim, ep_length)
+        self.action_space = MultiDiscrete([dim, dim])
+        self.observation_space = self.action_space
+        self.reset()
+
+
+class IdentityEnvMultiBinary(IdentityEnv):
+    def __init__(self, dim, ep_length=100):
+        """
+        Identity environment for testing purposes
+
+        :param dim: (int) the size of the dimensions you want to learn
+        :param ep_length: (int) the length of each episodes in timesteps
+        """
+        super(IdentityEnvMultiBinary, self).__init__(dim, ep_length)
+        self.action_space = MultiBinary(dim)
+        self.observation_space = self.action_space
+        self.reset()

torchy_baselines/common/save_util.py

@@ -0,0 +1,134 @@
+"""
+Save util taken from stable_baselines
+used to serialize data (class parameters) of model classes
+"""
+
+
+import json
+import base64
+import pickle
+import cloudpickle
+
+
+def is_json_serializable(item):
+    """
+    Test if an object is serializable into JSON
+
+    :param item: (object) The object to be tested for JSON serialization.
+    :return: (bool) True if object is JSON serializable, false otherwise.
+    """
+    # Try with try-except struct.
+    json_serializable = True
+    try:
+        _ = json.dumps(item)
+    except TypeError:
+        json_serializable = False
+    return json_serializable
+
+
+def data_to_json(data):
+    """
+    Turn data (class parameters) into a JSON string for storing
+
+    :param data: (Dict) Dictionary of class parameters to be
+        stored. Items that are not JSON serializable will be
+        pickled with Cloudpickle and stored as bytearray in
+        the JSON file
+    :return: (str) JSON string of the data serialized.
+    """
+    # First, check what elements can not be JSONfied,
+    # and turn them into byte-strings
+    serializable_data = {}
+    for data_key, data_item in data.items():
+        # See if object is JSON serializable
+        if is_json_serializable(data_item):
+            # All good, store as it is
+            serializable_data[data_key] = data_item
+        else:
+            # Not serializable, cloudpickle it into
+            # bytes and convert to base64 string for storing.
+            # Also store type of the class for consumption
+            # from other languages/humans, so we have an
+            # idea what was being stored.
+            base64_encoded = base64.b64encode(
+                cloudpickle.dumps(data_item)
+            ).decode()
+
+            # Use ":" to make sure we do
+            # not override these keys
+            # when we include variables of the object later
+            cloudpickle_serialization = {
+                ":type:": str(type(data_item)),
+                ":serialized:": base64_encoded
+            }
+
+            # Add first-level JSON-serializable items of the
+            # object for further details (but not deeper than this to
+            # avoid deep nesting).
+            # First we check that object has attributes (not all do,
+            # e.g. numpy scalars)
+            if hasattr(data_item, "__dict__") or isinstance(data_item, dict):
+                # Take elements from __dict__ for custom classes
+                item_generator = (
+                    data_item.items if isinstance(data_item, dict) else data_item.__dict__.items
+                )
+                for variable_name, variable_item in item_generator():
+                    # Check if serializable. If not, just include the
+                    # string-representation of the object.
+                    if is_json_serializable(variable_item):
+                        cloudpickle_serialization[variable_name] = variable_item
+                    else:
+                        cloudpickle_serialization[variable_name] = str(variable_item)
+
+            serializable_data[data_key] = cloudpickle_serialization
+    json_string = json.dumps(serializable_data, indent=4)
+    return json_string
+
+
+def json_to_data(json_string, custom_objects=None):
+    """
+    Turn JSON serialization of class-parameters back into dictionary.
+
+    :param json_string: (str) JSON serialization of the class-parameters
+        that should be loaded.
+    :param custom_objects: (dict) Dictionary of objects to replace
+        upon loading. If a variable is present in this dictionary as a
+        key, it will not be deserialized and the corresponding item
+        will be used instead. Similar to custom_objects in
+        `keras.models.load_model`. Useful when you have an object in
+        file that can not be deserialized.
+    :return: (dict) Loaded class parameters.
+    """
+    if custom_objects is not None and not isinstance(custom_objects, dict):
+        raise ValueError("custom_objects argument must be a dict or None")
+
+    json_dict = json.loads(json_string)
+    # This will be filled with deserialized data
+    return_data = {}
+    for data_key, data_item in json_dict.items():
+        if custom_objects is not None and data_key in custom_objects.keys():
+            # If item is provided in custom_objects, replace
+            # the one from JSON with the one in custom_objects
+            return_data[data_key] = custom_objects[data_key]
+        elif isinstance(data_item, dict) and ":serialized:" in data_item.keys():
+            # If item is dictionary with ":serialized:"
+            # key, this means it is serialized with cloudpickle.
+            serialization = data_item[":serialized:"]
+            # Try-except deserialization in case we run into
+            # errors. If so, we can tell bit more information to
+            # user.
+            try:
+                deserialized_object = cloudpickle.loads(
+                    base64.b64decode(serialization.encode())
+                )
+            except pickle.UnpicklingError:
+                raise RuntimeError(
+                    "Could not deserialize object {}. ".format(data_key) +
+                    "Consider using `custom_objects` argument to replace " +
+                    "this object."
+                )
+            return_data[data_key] = deserialized_object
+        else:
+            # Read as it is
+            return_data[data_key] = data_item
+    return return_data

torchy_baselines/ppo/ppo.py

@@ -6,6 +6,7 @@ import gym
 from gym import spaces
 import torch as th
 import torch.nn.functional as F
+
 # Check if tensorboard is available for pytorch
 try:
     from torch.utils.tensorboard import SummaryWriter
@@ -185,7 +186,6 @@ class PPO(BaseRLModel):
             clip_range_vf = self.clip_range_vf(self._current_progress)
             logger.logkv("clip_range_vf", clip_range_vf)
 
-
         for gradient_step in range(gradient_steps):
             approx_kl_divs = []
             # Sample replay buffer
@@ -219,7 +219,6 @@ class PPO(BaseRLModel):
                 # Value loss using the TD(gae_lambda) target
                 value_loss = F.mse_loss(return_batch, values_pred)
 
-
                 # Entropy loss favor exploration
                 entropy_loss = -th.mean(entropy)
 
@@ -234,7 +233,7 @@ class PPO(BaseRLModel):
                 approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy())
 
             if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl:
-                print("Early stopping at step {} due to reaching max kl: {:.2f}".format(it, np.mean(approx_kl_divs)))
+                print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step, np.mean(approx_kl_divs)))
                 break
 
         # print(explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
@@ -294,13 +293,39 @@ class PPO(BaseRLModel):
         return self
 
     def save(self, path):
-        if not path.endswith('.pth'):
-            path += '.pth'
-        th.save(self.policy.state_dict(), path)
+        """
+        saves all the params from init and pytorch params in a file for continous learning
 
-    def load(self, path, env=None, **_kwargs):
+        :param path: path to the file where the data should be safed
+        :return:
+        """
+
+        data = {
+            "gamma": self.gamma,
+            "n_steps": self.n_steps,
+            "vf_coef": self.vf_coef,
+            "ent_coef": self.ent_coef,
+            "max_grad_norm": self.max_grad_norm,
+            "learning_rate": self.learning_rate,
+            "gae_lambda": self.gae_lambda,
+            "n_epochs": self.n_epochs,
+            "clip_range": self.clip_range,
+            "clip_range_vf": self.clip_range_vf,
+            "batch_size": self.batch_size,
+            "target_kl": self.target_kl,
+            "tensorboard_log": self.tensorboard_log,
+            "policy_kwargs": self.policy_kwargs,
+            "policy": self.policy,
+
+        }
+
+        params_to_save = self.get_parameters()
+
+        self._save_to_file_zip(path, data=data, params=params_to_save)
+
+    """def load(self, path, env=None, **_kwargs):
         if not path.endswith('.pth'):
             path += '.pth'
         if env is not None:
             pass
-        self.policy.load_state_dict(th.load(path))
+        self.policy.load_state_dict(th.load(path))"""