diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index cd52748..5ea2904 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -16,12 +16,15 @@ New Features: - Added ``StopTrainingOnMaxEpisodes`` to callback collection (@xicocaio) - Added ``device`` keyword argument to ``BaseAlgorithm.load()`` (@liorcohen5) - Callbacks have access to rollout collection locals as in SB2. (@PartiallyTyped) +- Added ``get_parameters`` and ``set_parameters`` for accessing/setting parameters of the agent +- Added actor/critic loss logging for TD3. (@mloo3) Bug Fixes: ^^^^^^^^^^ - Fixed a bug where the environment was reset twice when using ``evaluate_policy`` - Fix logging of ``clip_fraction`` in PPO (@diditforlulz273) - Fixed a bug where cuda support was wrongly checked when passing the GPU index, e.g., ``device="cuda:0"`` (@liorcohen5) +- Fixed a bug when the random seed was not properly set on cuda when passing the GPU index Deprecations: ^^^^^^^^^^^^^ @@ -32,6 +35,14 @@ Others: - Fix type annotation of ``make_vec_env`` (@ManifoldFR) - Removed ``AlreadySteppingError`` and ``NotSteppingError`` that were not used - Fixed typos in SAC and TD3 +- Rename ``BaseClass.get_torch_variables`` -> ``BaseClass._get_torch_save_params`` and + ``BaseClass.excluded_save_params`` -> ``BaseClass._excluded_save_params`` +- Reorganized functions for clarity in ``BaseClass`` (save/load functions close to each other, private + functions at top) +- Clarified docstrings on what is saved and loaded to/from files +- Renamed saved items ``tensors`` to ``pytorch_variables`` for clarity +- Simplified ``save_to_zip_file`` function by removing duplicate code +- Store library version along with the saved models Documentation: ^^^^^^^^^^^^^^ @@ -403,4 +414,4 @@ And all the contributors: @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio -@diditforlulz273 @liorcohen5 @ManifoldFR +@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 5c87bca..5055000 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -221,181 +221,43 @@ class BaseAlgorithm(ABC): for optimizer in optimizers: update_learning_rate(optimizer, self.lr_schedule(self._current_progress_remaining)) - def get_env(self) -> Optional[VecEnv]: + def _excluded_save_params(self) -> List[str]: """ - Returns the current environment (can be None if not defined). + Returns the names of the parameters that should be excluded from being + saved by pickling. E.g. replay buffers are skipped by default + as they take up a lot of space. PyTorch variables should be excluded + with this so they can be stored with ``th.save``. - :return: (Optional[VecEnv]) The current environment + :return: (List[str]) List of parameters that should be excluded from being saved with pickle. """ - return self.env + return [ + "policy", + "device", + "env", + "eval_env", + "replay_buffer", + "rollout_buffer", + "_vec_normalize_env", + ] - def get_vec_normalize_env(self) -> Optional[VecNormalize]: + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: """ - Return the ``VecNormalize`` wrapper of the training env - if it exists. - :return: Optional[VecNormalize] The ``VecNormalize`` env. - """ - return self._vec_normalize_env + Get the name of the torch variables that will be saved with + PyTorch ``th.save``, ``th.load`` and ``state_dicts`` instead of the default + pickling strategy. This is to handle device placement correctly. - def set_env(self, env: GymEnv) -> None: - """ - Checks the validity of the environment, and if it is coherent, set it as the current environment. - Furthermore wrap any non vectorized env into a vectorized - checked parameters: - - observation_space - - action_space - - :param env: The environment for learning a policy - """ - check_for_correct_spaces(env, self.observation_space, self.action_space) - # it must be coherent now - # if it is not a VecEnv, make it a VecEnv - env = self._wrap_env(env) - - self.n_envs = env.num_envs - self.env = env - - def get_torch_variables(self) -> Tuple[List[str], List[str]]: - """ - Get the name of the torch variables that will be saved. - ``th.save`` and ``th.load`` will be used with the right device - instead of the default pickling strategy. + Names can point to specific variables under classes, e.g. + "policy.optimizer" would point to ``optimizer`` object of ``self.policy`` + if this object. :return: (Tuple[List[str], List[str]]) - name of the variables with state dicts to save, name of additional torch tensors, + List of Torch variables whose state dicts to save (e.g. th.nn.Modules), + and list of other Torch variables to store with ``th.save``. """ state_dicts = ["policy"] return state_dicts, [] - @abstractmethod - def learn( - self, - total_timesteps: int, - callback: MaybeCallback = None, - log_interval: int = 100, - tb_log_name: str = "run", - eval_env: Optional[GymEnv] = None, - eval_freq: int = -1, - n_eval_episodes: int = 5, - eval_log_path: Optional[str] = None, - reset_num_timesteps: bool = True, - ) -> "BaseAlgorithm": - """ - Return a trained model. - - :param total_timesteps: (int) The total number of samples (env steps) to train on - :param callback: (MaybeCallback) callback(s) called at every step with state of the algorithm. - :param log_interval: (int) The number of timesteps before logging. - :param tb_log_name: (str) the name of the run for TensorBoard logging - :param eval_env: (gym.Env) Environment that will be used to evaluate the agent - :param eval_freq: (int) Evaluate the agent every ``eval_freq`` timesteps (this may vary a little) - :param n_eval_episodes: (int) Number of episode to evaluate the agent - :param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved - :param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging) - :return: (BaseAlgorithm) the trained model - """ - - def predict( - self, - observation: np.ndarray, - state: Optional[np.ndarray] = None, - mask: Optional[np.ndarray] = None, - deterministic: bool = False, - ) -> Tuple[np.ndarray, Optional[np.ndarray]]: - """ - Get the model's action(s) from an observation - - :param observation: (np.ndarray) the input observation - :param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies) - :param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies) - :param deterministic: (bool) Whether or not to return deterministic actions. - :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state - (used in recurrent policies) - """ - return self.policy.predict(observation, state, mask, deterministic) - - @classmethod - def load( - cls, load_path: str, env: Optional[GymEnv] = None, device: Union[th.device, str] = "auto", **kwargs - ) -> "BaseAlgorithm": - """ - Load the model from a zip-file - - :param load_path: the location of the saved data - :param env: the new environment to run the loaded model on - (can be None if you only need prediction from a trained model) has priority over any saved environment - :param device: (Union[th.device, str]) Device on which the code should run. - :param kwargs: extra arguments to change the model when loading - """ - data, params, tensors = load_from_zip_file(load_path, device=device) - - if "policy_kwargs" in data: - for arg_to_remove in ["device"]: - if arg_to_remove in data["policy_kwargs"]: - del data["policy_kwargs"][arg_to_remove] - - if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]: - raise ValueError( - f"The specified policy kwargs do not equal the stored policy kwargs." - f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}" - ) - - # check if observation space and action space are part of the saved parameters - if "observation_space" not in data or "action_space" not in data: - raise KeyError("The observation_space and action_space were not given, can't verify new environments") - # check if given env is valid - if env is not None: - check_for_correct_spaces(env, data["observation_space"], data["action_space"]) - # if no new env was given use stored env if possible - if env is None and "env" in data: - env = data["env"] - - # noinspection PyArgumentList - model = cls( - policy=data["policy_class"], - env=env, - device=device, - _init_setup_model=False, # pytype: disable=not-instantiable,wrong-keyword-args - ) - - # load parameters - model.__dict__.update(data) - model.__dict__.update(kwargs) - model._setup_model() - - # put state_dicts back in place - for name in params: - attr = recursive_getattr(model, name) - attr.load_state_dict(params[name]) - - # put tensors back in place - if tensors is not None: - for name in tensors: - recursive_setattr(model, name, tensors[name]) - - # Sample gSDE exploration matrix, so it uses the right device - # see issue #44 - if model.use_sde: - model.policy.reset_noise() # pytype: disable=attribute-error - return model - - def set_random_seed(self, seed: Optional[int] = None) -> None: - """ - Set the seed of the pseudo-random generators - (python, numpy, pytorch, gym, action_space) - - :param seed: (int) - """ - if seed is None: - return - set_random_seed(seed, using_cuda=self.device == th.device("cuda")) - self.action_space.seed(seed) - if self.env is not None: - self.env.seed(seed) - if self.eval_env is not None: - self.eval_env.seed(seed) - def _init_callback( self, callback: MaybeCallback, @@ -513,14 +375,254 @@ class BaseAlgorithm(ABC): if maybe_is_success is not None and dones[idx]: self.ep_success_buffer.append(maybe_is_success) - def excluded_save_params(self) -> List[str]: + def get_env(self) -> Optional[VecEnv]: """ - Returns the names of the parameters that should be excluded by default - when saving the model. + Returns the current environment (can be None if not defined). - :return: ([str]) List of parameters that should be excluded from save + :return: (Optional[VecEnv]) The current environment """ - return ["policy", "device", "env", "eval_env", "replay_buffer", "rollout_buffer", "_vec_normalize_env"] + return self.env + + def get_vec_normalize_env(self) -> Optional[VecNormalize]: + """ + Return the ``VecNormalize`` wrapper of the training env + if it exists. + :return: Optional[VecNormalize] The ``VecNormalize`` env. + """ + return self._vec_normalize_env + + def set_env(self, env: GymEnv) -> None: + """ + Checks the validity of the environment, and if it is coherent, set it as the current environment. + Furthermore wrap any non vectorized env into a vectorized + checked parameters: + - observation_space + - action_space + + :param env: The environment for learning a policy + """ + check_for_correct_spaces(env, self.observation_space, self.action_space) + # it must be coherent now + # if it is not a VecEnv, make it a VecEnv + env = self._wrap_env(env) + + self.n_envs = env.num_envs + self.env = env + + @abstractmethod + def learn( + self, + total_timesteps: int, + callback: MaybeCallback = None, + log_interval: int = 100, + tb_log_name: str = "run", + eval_env: Optional[GymEnv] = None, + eval_freq: int = -1, + n_eval_episodes: int = 5, + eval_log_path: Optional[str] = None, + reset_num_timesteps: bool = True, + ) -> "BaseAlgorithm": + """ + Return a trained model. + + :param total_timesteps: (int) The total number of samples (env steps) to train on + :param callback: (MaybeCallback) callback(s) called at every step with state of the algorithm. + :param log_interval: (int) The number of timesteps before logging. + :param tb_log_name: (str) the name of the run for TensorBoard logging + :param eval_env: (gym.Env) Environment that will be used to evaluate the agent + :param eval_freq: (int) Evaluate the agent every ``eval_freq`` timesteps (this may vary a little) + :param n_eval_episodes: (int) Number of episode to evaluate the agent + :param eval_log_path: (Optional[str]) Path to a folder where the evaluations will be saved + :param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging) + :return: (BaseAlgorithm) the trained model + """ + + def predict( + self, + observation: np.ndarray, + state: Optional[np.ndarray] = None, + mask: Optional[np.ndarray] = None, + deterministic: bool = False, + ) -> Tuple[np.ndarray, Optional[np.ndarray]]: + """ + Get the model's action(s) from an observation + + :param observation: (np.ndarray) the input observation + :param state: (Optional[np.ndarray]) The last states (can be None, used in recurrent policies) + :param mask: (Optional[np.ndarray]) The last masks (can be None, used in recurrent policies) + :param deterministic: (bool) Whether or not to return deterministic actions. + :return: (Tuple[np.ndarray, Optional[np.ndarray]]) the model's action and the next state + (used in recurrent policies) + """ + return self.policy.predict(observation, state, mask, deterministic) + + def set_random_seed(self, seed: Optional[int] = None) -> None: + """ + Set the seed of the pseudo-random generators + (python, numpy, pytorch, gym, action_space) + + :param seed: (int) + """ + if seed is None: + return + set_random_seed(seed, using_cuda=self.device.type == th.device("cuda").type) + self.action_space.seed(seed) + if self.env is not None: + self.env.seed(seed) + if self.eval_env is not None: + self.eval_env.seed(seed) + + def set_parameters( + self, + load_path_or_dict: Union[str, Dict[str, Dict]], + exact_match: bool = True, + device: Union[th.device, str] = "auto", + ): + """ + Load parameters from a given zip-file or a nested dictionary containing parameters for + different modules (see ``get_parameters``). + + :param load_path_or_iter: Location of the saved data (path or file-like, see ``save``), or a nested + dictionary containing nn.Module parameters used by the policy. The dictionary maps + object names to a state-dictionary returned by ``torch.nn.Module.state_dict()``. + :param exact_match: If True, the given parameters should include parameters for each + module and each of their parameters, otherwise raises an Exception. If set to False, this + can be used to update only specific parameters. + :param device: (Union[th.device, str]) Device on which the code should run. + """ + params = None + if isinstance(load_path_or_dict, dict): + params = load_path_or_dict + else: + _, params, _ = load_from_zip_file(load_path_or_dict, device=device) + + # Keep track which objects were updated. + # `_get_torch_save_params` returns [params, other_pytorch_variables]. + # We are only interested in former here. + objects_needing_update = set(self._get_torch_save_params()[0]) + updated_objects = set() + + for name in params: + attr = None + try: + attr = recursive_getattr(self, name) + except Exception: + # What errors recursive_getattr could throw? KeyError, but + # possible something else too (e.g. if key is an int?). + # Catch anything for now. + raise ValueError(f"Key {name} is an invalid object name.") + + if isinstance(attr, th.optim.Optimizer): + # Optimizers do not support "strict" keyword... + # Seems like they will just replace the whole + # optimizer state with the given one. + # On top of this, optimizer state-dict + # seems to change (e.g. first ``optim.step()``), + # which makes comparing state dictionary keys + # invalid (there is also a nesting of dictionaries + # with lists with dictionaries with ...), adding to the + # mess. + # + # TL;DR: We might not be able to reliably say + # if given state-dict is missing keys. + # + # Solution: Just load the state-dict as is, and trust + # the user has provided a sensible state dictionary. + attr.load_state_dict(params[name]) + else: + # Assume attr is th.nn.Module + attr.load_state_dict(params[name], strict=exact_match) + updated_objects.add(name) + + if exact_match and updated_objects != objects_needing_update: + raise ValueError( + "Names of parameters do not match agents' parameters: " + f"expected {objects_needing_update}, got {updated_objects}" + ) + + @classmethod + def load( + cls, + path: Union[str, pathlib.Path, io.BufferedIOBase], + env: Optional[GymEnv] = None, + device: Union[th.device, str] = "auto", + **kwargs, + ) -> "BaseAlgorithm": + """ + Load the model from a zip-file + + :param path: (Union[str, pathlib.Path, io.BufferedIOBase]) path to the file (or a file-like) where to + load the agent from + :param env: the new environment to run the loaded model on + (can be None if you only need prediction from a trained model) has priority over any saved environment + :param device: (Union[th.device, str]) Device on which the code should run. + :param kwargs: extra arguments to change the model when loading + """ + data, params, pytorch_variables = load_from_zip_file(path, device=device) + + # Remove stored device information and replace with ours + if "policy_kwargs" in data: + if "device" in "policy_kwargs": + del data["policy_kwargs"]["device"] + + if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]: + raise ValueError( + f"The specified policy kwargs do not equal the stored policy kwargs." + f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}" + ) + + if "observation_space" not in data or "action_space" not in data: + raise KeyError("The observation_space and action_space were not given, can't verify new environments") + + if env is not None: + # Check if given env is valid + check_for_correct_spaces(env, data["observation_space"], data["action_space"]) + else: + # Use stored env, if one exists. If not, continue as is (can be used for predict) + if "env" in data: + env = data["env"] + + # noinspection PyArgumentList + model = cls( + policy=data["policy_class"], + env=env, + device=device, + _init_setup_model=False, # pytype: disable=not-instantiable,wrong-keyword-args + ) + + # load parameters + model.__dict__.update(data) + model.__dict__.update(kwargs) + model._setup_model() + + # put state_dicts back in place + model.set_parameters(params, exact_match=True, device=device) + + # put other pytorch variables back in place + if pytorch_variables is not None: + for name in pytorch_variables: + recursive_setattr(model, name, pytorch_variables[name]) + + # Sample gSDE exploration matrix, so it uses the right device + # see issue #44 + if model.use_sde: + model.policy.reset_noise() # pytype: disable=attribute-error + return model + + def get_parameters(self): + """ + Return the parameters of the agent. This includes parameters from different networks, e.g. + critics (value functions) and policies (pi functions). + + :return: (Dict[str, Dict]) Mapping of from names of the objects to PyTorch state-dicts. + """ + state_dicts_names, _ = self._get_torch_save_params() + params = {} + for name in state_dicts_names: + attr = recursive_getattr(self, name) + # Retrieve state dict + params[name] = attr.state_dict() + return params def save( self, @@ -532,46 +634,42 @@ class BaseAlgorithm(ABC): Save all the attributes of the object and the model parameters in a zip-file. :param (Union[str, pathlib.Path, io.BufferedIOBase]): path to the file where the rl agent should be saved - :param exclude: name of parameters that should be excluded in addition to the default one + :param exclude: name of parameters that should be excluded in addition to the default ones :param include: name of parameters that might be excluded but should be included anyway """ - # copy parameter list so we don't mutate the original dict + # Copy parameter list so we don't mutate the original dict data = self.__dict__.copy() # Exclude is union of specified parameters (if any) and standard exclusions if exclude is None: exclude = [] - exclude = set(exclude).union(self.excluded_save_params()) + exclude = set(exclude).union(self._excluded_save_params()) # Do not exclude params if they are specifically included if include is not None: exclude = exclude.difference(include) - state_dicts_names, tensors_names = self.get_torch_variables() - # any params that are in the save vars must not be saved by data - torch_variables = state_dicts_names + tensors_names - for torch_var in torch_variables: - # we need to get only the name of the top most module as we'll remove that + state_dicts_names, torch_variable_names = self._get_torch_save_params() + all_pytorch_variables = state_dicts_names + torch_variable_names + for torch_var in all_pytorch_variables: + # We need to get only the name of the top most module as we'll remove that var_name = torch_var.split(".")[0] + # Any params that are in the save vars must not be saved by data exclude.add(var_name) # Remove parameter entries of parameters which are to be excluded for param_name in exclude: data.pop(param_name, None) - # Build dict of tensor variables - tensors = None - if tensors_names is not None: - tensors = {} - for name in tensors_names: + # Build dict of torch variables + pytorch_variables = None + if torch_variable_names is not None: + pytorch_variables = {} + for name in torch_variable_names: attr = recursive_getattr(self, name) - tensors[name] = attr + pytorch_variables[name] = attr # Build dict of state_dicts - params_to_save = {} - for name in state_dicts_names: - attr = recursive_getattr(self, name) - # Retrieve state dict - params_to_save[name] = attr.state_dict() + params_to_save = self.get_parameters() - save_to_zip_file(path, data=data, params=params_to_save, tensors=tensors) + save_to_zip_file(path, data=data, params=params_to_save, pytorch_variables=pytorch_variables) diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index 721bf40..e8e5bf5 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -240,10 +240,7 @@ class OnPolicyAlgorithm(BaseAlgorithm): return self - def get_torch_variables(self) -> Tuple[List[str], List[str]]: - """ - cf base class - """ + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: state_dicts = ["policy", "policy.optimizer"] return state_dicts, [] diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 0decdc9..7260f11 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -16,6 +16,7 @@ from typing import Any, Dict, Optional, Tuple, Union import cloudpickle import torch as th +import stable_baselines3 from stable_baselines3.common.type_aliases import TensorDict from stable_baselines3.common.utils import get_device @@ -284,21 +285,20 @@ def save_to_zip_file( save_path: Union[str, pathlib.Path, io.BufferedIOBase], data: Dict[str, Any] = None, params: Dict[str, Any] = None, - tensors: Dict[str, Any] = None, + pytorch_variables: Dict[str, Any] = None, verbose=0, ) -> None: """ - Save a model to a zip archive. + Save model data to a zip archive. :param save_path: (Union[str, pathlib.Path, io.BufferedIOBase]) Where to store the model. if save_path is a str or pathlib.Path ensures that the path actually exists. - :param data: Class parameters being stored. + :param data: Class parameters being stored (non-PyTorch variables) :param params: Model parameters being stored expected to contain an entry for every state_dict with its name and the state_dict. - :param tensors: Extra tensor variables expected to contain name and value of tensors + :param pytorch_variables: Other PyTorch variables expected to contain name and value of the variable. :param verbose: (int) Verbosity level, 0 means only warnings, 2 means debug information """ - save_path = open_path(save_path, "w", verbose=0, suffix="zip") # data/params can be None, so do not # try to serialize them blindly @@ -310,13 +310,15 @@ def save_to_zip_file( # Do not try to save "None" elements if data is not None: archive.writestr("data", serialized_data) - if tensors is not None: - with archive.open("tensors.pth", mode="w") as tensors_file: - th.save(tensors, tensors_file) + if pytorch_variables is not None: + with archive.open("pytorch_variables.pth", mode="w") as pytorch_variables_file: + th.save(pytorch_variables, pytorch_variables_file) if params is not None: for file_name, dict_ in params.items(): with archive.open(file_name + ".pth", mode="w") as param_file: th.save(dict_, param_file) + # Save metadata: library version when file was saved + archive.writestr("_stable_baselines3_version", stable_baselines3.__version__) def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj, verbose=0) -> None: @@ -362,8 +364,8 @@ def load_from_zip_file( :param load_data: Whether we should load and return data (class parameters). Mainly used by 'load_parameters' to only load model parameters (weights) :param device: (Union[th.device, str]) Device on which the code should run. - :return: (dict),(dict),(dict) Class parameters, model state_dicts (dict of state_dict) - and dict of extra tensors + :return: (dict),(dict),(dict) Class parameters, model state_dicts (aka "params", dict of state_dict) + and dict of pytorch variables """ load_path = open_path(load_path, "r", verbose=verbose, suffix="zip") @@ -378,44 +380,38 @@ def load_from_zip_file( # zip archive, assume they were stored # as None (_save_to_file_zip allows this). data = None - tensors = None + pytorch_variables = None params = {} if "data" in namelist and load_data: - # Load class parameters and convert to string + # Load class parameters that are stored + # with either JSON or pickle (not PyTorch variables). json_data = archive.read("data").decode() data = json_to_data(json_data) - if "tensors.pth" in namelist and load_data: - # Load extra tensors - with archive.open("tensors.pth", mode="r") as tensor_file: - # File has to be seekable, but opt_param_file is not, so load in BytesIO first + # Check for all .pth files and load them using th.load. + # "pytorch_variables.pth" stores PyTorch variables, and any other .pth + # files store state_dicts of variables with custom names (e.g. policy, policy.optimizer) + pth_files = [file_name for file_name in namelist if os.path.splitext(file_name)[1] == ".pth"] + for file_path in pth_files: + with archive.open(file_path, mode="r") as param_file: + # File has to be seekable, but param_file is not, so load in BytesIO first # fixed in python >= 3.7 file_content = io.BytesIO() - file_content.write(tensor_file.read()) + file_content.write(param_file.read()) # go to start of file file_content.seek(0) - # load the parameters with the right ``map_location`` - tensors = th.load(file_content, map_location=device) - - # check for all other .pth files - other_files = [ - file_name for file_name in namelist if os.path.splitext(file_name)[1] == ".pth" and file_name != "tensors.pth" - ] - # if there are any other files which end with .pth and aren't "params.pth" - # assume that they each are optimizer parameters - if len(other_files) > 0: - for file_path in other_files: - with archive.open(file_path, mode="r") as opt_param_file: - # File has to be seekable, but opt_param_file is not, so load in BytesIO first - # fixed in python >= 3.7 - file_content = io.BytesIO() - file_content.write(opt_param_file.read()) - # go to start of file - file_content.seek(0) - # load the parameters with the right ``map_location`` - params[os.path.splitext(file_path)[0]] = th.load(file_content, map_location=device) + # Load the parameters with the right ``map_location``. + # Remove ".pth" ending with splitext + th_object = th.load(file_content, map_location=device) + if file_path == "pytorch_variables.pth": + # PyTorch variables (not state_dicts) + pytorch_variables = th_object + else: + # State dicts. Store into params dictionary + # with same name as in .zip file (without .pth) + params[os.path.splitext(file_path)[0]] = th_object except zipfile.BadZipFile: # load_path wasn't a zip file raise ValueError(f"Error: the file {load_path} wasn't a zip-file") - return data, params, tensors + return data, params, pytorch_variables diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index b030b43..3583c70 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -236,20 +236,10 @@ class DQN(OffPolicyAlgorithm): reset_num_timesteps=reset_num_timesteps, ) - def excluded_save_params(self) -> List[str]: - """ - Returns the names of the parameters that should be excluded by default - when saving the model. + def _excluded_save_params(self) -> List[str]: + return super(DQN, self)._excluded_save_params() + ["q_net", "q_net_target"] - :return: (List[str]) List of parameters that should be excluded from save - """ - # Exclude aliases - return super(DQN, self).excluded_save_params() + ["q_net", "q_net_target"] - - def get_torch_variables(self) -> Tuple[List[str], List[str]]: - """ - cf base class - """ + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: state_dicts = ["policy", "policy.optimizer"] return state_dicts, [] diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index c043171..b12aaf0 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -548,32 +548,20 @@ class SAC(OffPolicyAlgorithm): reset_num_timesteps=reset_num_timesteps, ) - def excluded_save_params(self) -> List[str]: - """ - Returns the names of the parameters that should be excluded by default - when saving the model. + def _excluded_save_params(self) -> List[str]: + return super(SAC, self)._excluded_save_params() + ["actor", "critic", "critic_target"] - :return: (List[str]) List of parameters that should be excluded from save - """ - # Exclude aliases - return super(SAC, self).excluded_save_params() + ["actor", "critic", "critic_target"] - - def get_torch_variables(self) -> Tuple[List[str], List[str]]: - """ - cf base class - """ + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: state_dicts = ["policy", "actor.optimizer", "critic.optimizer"] - saved_tensors = ["log_ent_coef"] + saved_pytorch_variables = ["log_ent_coef"] if self.ent_coef_optimizer is not None: state_dicts.append("ent_coef_optimizer") else: - saved_tensors.append("ent_coef_tensor") - + saved_pytorch_variables.append("ent_coef_tensor") if self.use_cql: if self.alpha_coef == "auto": state_dicts.append("alpha_optimizer") - saved_tensors.append("log_alpha") + saved_pytorch_variables.append("log_alpha") else: - saved_tensors.append("alpha_coef_tensor") - - return state_dicts, saved_tensors + saved_pytorch_variables.append("alpha_coef_tensor") + return state_dicts, saved_pytorch_variables diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 7738b0b..4b51c4b 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -1,5 +1,6 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union +import numpy as np import torch as th from torch.nn import functional as F @@ -135,6 +136,8 @@ class TD3(OffPolicyAlgorithm): # Update learning rate according to lr schedule self._update_learning_rate([self.actor.optimizer, self.critic.optimizer]) + actor_losses, critic_losses = [], [] + for gradient_step in range(gradient_steps): # Sample replay buffer @@ -156,6 +159,7 @@ class TD3(OffPolicyAlgorithm): # Compute critic loss critic_loss = sum([F.mse_loss(current_q, target_q) for current_q in current_q_estimates]) + critic_losses.append(critic_loss.item()) # Optimize the critics self.critic.optimizer.zero_grad() @@ -166,6 +170,7 @@ class TD3(OffPolicyAlgorithm): if gradient_step % self.policy_delay == 0: # Compute actor loss actor_loss = -self.critic.q1_forward(replay_data.observations, self.actor(replay_data.observations)).mean() + actor_losses.append(actor_loss.item()) # Optimize the actor self.actor.optimizer.zero_grad() @@ -177,6 +182,8 @@ class TD3(OffPolicyAlgorithm): self._n_updates += gradient_steps logger.record("train/n_updates", self._n_updates, exclude="tensorboard") + logger.record("train/actor_loss", np.mean(actor_losses)) + logger.record("train/critic_loss", np.mean(critic_losses)) def learn( self, @@ -203,19 +210,9 @@ class TD3(OffPolicyAlgorithm): reset_num_timesteps=reset_num_timesteps, ) - def excluded_save_params(self) -> List[str]: - """ - Returns the names of the parameters that should be excluded by default - when saving the model. + def _excluded_save_params(self) -> List[str]: + return super(TD3, self)._excluded_save_params() + ["actor", "critic", "actor_target", "critic_target"] - :return: (List[str]) List of parameters that should be excluded from save - """ - # Exclude aliases - return super(TD3, self).excluded_save_params() + ["actor", "critic", "actor_target", "critic_target"] - - def get_torch_variables(self) -> Tuple[List[str], List[str]]: - """ - cf base class - """ + def _get_torch_save_params(self) -> Tuple[List[str], List[str]]: state_dicts = ["policy", "actor.optimizer", "critic.optimizer"] return state_dicts, [] diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 8ebeabd..dde9d2b 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -2,6 +2,7 @@ import io import os import pathlib import warnings +from collections import OrderedDict from copy import deepcopy import gym @@ -33,7 +34,7 @@ def select_env(model_class: BaseAlgorithm) -> gym.Env: def test_save_load(tmp_path, model_class): """ Test if 'save' and 'load' saves and loads model correctly - and if 'load_parameters' and 'get_policy_parameters' work correctly + and if 'get_parameters' and 'set_parameters' and work correctly. ''warning does not test function of optimizer parameter load @@ -49,19 +50,73 @@ def test_save_load(tmp_path, model_class): env.reset() observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0) - # Get dictionary of current parameters - params = deepcopy(model.policy.state_dict()) + # Get parameters of different objects + # deepcopy to avoid referencing to tensors we are about to modify + original_params = deepcopy(model.get_parameters()) - # Modify all parameters to be random values - random_params = dict((param_name, th.rand_like(param)) for param_name, param in params.items()) + # Test different error cases of set_parameters. + # Test that invalid object names throw errors + invalid_object_params = deepcopy(original_params) + invalid_object_params["I_should_not_be_a_valid_object"] = "and_I_am_an_invalid_tensor" + with pytest.raises(ValueError): + model.set_parameters(invalid_object_params, exact_match=True) + with pytest.raises(ValueError): + model.set_parameters(invalid_object_params, exact_match=False) + + # Test that exact_match catches when something was missed. + missing_object_params = dict((k, v) for k, v in list(original_params.items())[:-1]) + with pytest.raises(ValueError): + model.set_parameters(missing_object_params, exact_match=True) + + # Test that exact_match catches when something inside state-dict + # is missing but we have exact_match. + missing_state_dict_tensor_params = {} + for object_name in original_params: + object_params = {} + missing_state_dict_tensor_params[object_name] = object_params + # Skip last item in state-dict + for k, v in list(original_params[object_name].items())[:-1]: + object_params[k] = v + with pytest.raises(RuntimeError): + # PyTorch load_state_dict throws RuntimeError if strict but + # invalid state-dict. + model.set_parameters(missing_state_dict_tensor_params, exact_match=True) + + # Test that parameters do indeed change. + random_params = {} + for object_name, params in original_params.items(): + # Do not randomize optimizer parameters (custom layout) + if "optim" in object_name: + random_params[object_name] = params + else: + # Again, skip the last item in state-dict + random_params[object_name] = OrderedDict( + (param_name, th.rand_like(param)) for param_name, param in list(params.items())[:-1] + ) # Update model parameters with the new random values - model.policy.load_state_dict(random_params) + model.set_parameters(random_params, exact_match=False) - new_params = model.policy.state_dict() - # Check that all params are different now - for k in params: - assert not th.allclose(params[k], new_params[k]), "Parameters did not change as expected." + new_params = model.get_parameters() + # Check that all params except the final item in each state-dict are different. + for object_name in original_params: + # Skip optimizers (no valid comparison with just th.allclose) + if "optim" in object_name: + continue + # state-dicts use ordered dictionaries, so key order + # is guaranteed. + last_key = list(original_params[object_name].keys())[-1] + for k in original_params[object_name]: + if k == last_key: + # Should be same as before + assert th.allclose( + original_params[object_name][k], new_params[object_name][k] + ), "Parameter changed despite not included in the loaded parameters." + else: + # Should be different + assert not th.allclose( + original_params[object_name][k], new_params[object_name][k] + ), "Parameters did not change as expected." params = new_params @@ -81,14 +136,18 @@ def test_save_load(tmp_path, model_class): assert model.policy.device.type == get_device(device).type # check if params are still the same after load - new_params = model.policy.state_dict() + new_params = model.get_parameters() # Check that all params are the same as before save load procedure now - for key in params: - assert new_params[key].device.type == get_device(device).type - assert th.allclose( - params[key].to("cpu"), new_params[key].to("cpu") - ), "Model parameters not the same after save and load." + for object_name in new_params: + # Skip optimizers (no valid comparison with just th.allclose) + if "optim" in object_name: + continue + for key in params[object_name]: + assert new_params[object_name][key].device.type == get_device(device).type + assert th.allclose( + params[object_name][key].to("cpu"), new_params[object_name][key].to("cpu") + ), "Model parameters not the same after save and load." # check if model still selects the same actions new_selected_actions, _ = model.predict(observations, deterministic=True)