From cc744a48b5047edca8c02268d62f08e2e8f0f582 Mon Sep 17 00:00:00 2001 From: Noah Dormann Date: Tue, 12 Nov 2019 17:03:57 +0100 Subject: [PATCH] first save and load features --- tests/test_run.py | 10 ++ tests/test_save_load.py | 50 +++++++++ torchy_baselines/a2c/a2c.py | 12 +++ torchy_baselines/common/base_class.py | 132 ++++++++++++++++++++--- torchy_baselines/common/identity_env.py | 105 +++++++++++++++++++ torchy_baselines/common/save_util.py | 134 ++++++++++++++++++++++++ torchy_baselines/ppo/ppo.py | 41 ++++++-- 7 files changed, 459 insertions(+), 25 deletions(-) create mode 100644 tests/test_save_load.py create mode 100644 torchy_baselines/common/identity_env.py create mode 100644 torchy_baselines/common/save_util.py diff --git a/tests/test_run.py b/tests/test_run.py index 32a4b30..b578824 100644 --- a/tests/test_run.py +++ b/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) diff --git a/tests/test_save_load.py b/tests/test_save_load.py new file mode 100644 index 0000000..ca3f637 --- /dev/null +++ b/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") diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 6ee6f4a..57faabd 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/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)) \ No newline at end of file diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 6adc45c..7daffe2 100644 --- a/torchy_baselines/common/base_class.py +++ b/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) diff --git a/torchy_baselines/common/identity_env.py b/torchy_baselines/common/identity_env.py new file mode 100644 index 0000000..d815220 --- /dev/null +++ b/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() diff --git a/torchy_baselines/common/save_util.py b/torchy_baselines/common/save_util.py new file mode 100644 index 0000000..30dd0b2 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 13a1634..e6467f3 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/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))"""