Get/set parameters and review of saving and loading (#138)

* Update comments and docstrings

* Rename get_torch_variables to private and update docs

* Clarify documentation on data, params and tensors

* Make excluded_save_params private and update docs

* Update get_torch_variable_names to get_torch_save_params for description

* Simplify saving code and update docs on params vs tensors

* Rename saved item tensors to pytorch_variables for clarity

* Reformat

* Fix a typo

* Add get/set_parameters, update tests accordingly

* Use f-strings for formatting

* Fix load docstring

* Reorganize functions in BaseClass

* Update changelog

* Add library version to the stored models

* Actually run isort this time

* Fix flake8 complaints and also fix testing code

* Fix isort

* ...and black

* Fix set_random_seed

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
Co-authored-by: Antonin Raffin <antonin.raffin@dlr.de>

docs/misc/changelog.rst

@@ -16,6 +16,7 @@ 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:
@@ -23,6 +24,7 @@ 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:
 ^^^^^^^^^^^^^
@@ -33,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:
 ^^^^^^^^^^^^^^

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
+        Get the name of the torch variables that will be saved with
-        if it exists.
+        PyTorch ``th.save``, ``th.load`` and ``state_dicts`` instead of the default
-        :return: Optional[VecNormalize] The ``VecNormalize`` env.
+        pickling strategy. This is to handle device placement correctly.
-        """
-        return self._vec_normalize_env
 
-    def set_env(self, env: GymEnv) -> None:
+        Names can point to specific variables under classes, e.g.
-        """
+        "policy.optimizer" would point to ``optimizer`` object of ``self.policy``
-        Checks the validity of the environment, and if it is coherent, set it as the current environment.
+        if this object.
-        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.
 
         :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
+        Returns the current environment (can be None if not defined).
-        when saving the model.
 
-        :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()
+        state_dicts_names, torch_variable_names = self._get_torch_save_params()
-        # any params that are in the save vars must not be saved by data
+        all_pytorch_variables = state_dicts_names + torch_variable_names
-        torch_variables = state_dicts_names + tensors_names
+        for torch_var in all_pytorch_variables:
-        for torch_var in torch_variables:
+            # We need to get only the name of the top most module as we'll remove that
-            # 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
+        # Build dict of torch variables
-        tensors = None
+        pytorch_variables = None
-        if tensors_names is not None:
+        if torch_variable_names is not None:
-            tensors = {}
+            pytorch_variables = {}
-            for name in tensors_names:
+            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 = {}
+        params_to_save = self.get_parameters()
-        for name in state_dicts_names:
-            attr = recursive_getattr(self, name)
-            # Retrieve state dict
-            params_to_save[name] = attr.state_dict()
 
-        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)

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]]:
+    def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
-        """
-        cf base class
-        """
         state_dicts = ["policy", "policy.optimizer"]
 
         return state_dicts, []

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:
+        if pytorch_variables is not None:
-            with archive.open("tensors.pth", mode="w") as tensors_file:
+            with archive.open("pytorch_variables.pth", mode="w") as pytorch_variables_file:
-                th.save(tensors, tensors_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)
+    :return: (dict),(dict),(dict) Class parameters, model state_dicts (aka "params", dict of state_dict)
-        and dict of extra tensors
+        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:
+            # Check for all .pth files and load them using th.load.
-                # Load extra tensors
+            # "pytorch_variables.pth" stores PyTorch variables, and any other .pth
-                with archive.open("tensors.pth", mode="r") as tensor_file:
+            # files store state_dicts of variables with custom names (e.g. policy, policy.optimizer)
-                    # File has to be seekable, but opt_param_file is not, so load in BytesIO first
+            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``
+                    # Load the parameters with the right ``map_location``.
-                    tensors = th.load(file_content, map_location=device)
+                    # Remove ".pth" ending with splitext
-
+                    th_object = th.load(file_content, map_location=device)
-            # check for all other .pth files
+                    if file_path == "pytorch_variables.pth":
-            other_files = [
+                        # PyTorch variables (not state_dicts)
-                file_name for file_name in namelist if os.path.splitext(file_name)[1] == ".pth" and file_name != "tensors.pth"
+                        pytorch_variables = th_object
-            ]
+                    else:
-            # if there are any other files which end with .pth and aren't "params.pth"
+                        # State dicts. Store into params dictionary
-            # assume that they each are optimizer parameters
+                        # with same name as in .zip file (without .pth)
-            if len(other_files) > 0:
+                        params[os.path.splitext(file_path)[0]] = th_object
-                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)
     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

stable_baselines3/dqn/dqn.py

@@ -231,20 +231,10 @@ class DQN(OffPolicyAlgorithm):
             reset_num_timesteps=reset_num_timesteps,
         )
 
-    def excluded_save_params(self) -> List[str]:
+    def _excluded_save_params(self) -> List[str]:
-        """
+        return super(DQN, self)._excluded_save_params() + ["q_net", "q_net_target"]
-        Returns the names of the parameters that should be excluded by default
-        when saving the model.
 
-        :return: (List[str]) List of parameters that should be excluded from save
+    def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
-        """
-        # 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
-        """
         state_dicts = ["policy", "policy.optimizer"]
 
         return state_dicts, []

stable_baselines3/sac/sac.py

@@ -293,24 +293,14 @@ class SAC(OffPolicyAlgorithm):
             reset_num_timesteps=reset_num_timesteps,
         )
 
-    def excluded_save_params(self) -> List[str]:
+    def _excluded_save_params(self) -> List[str]:
-        """
+        return super(SAC, self)._excluded_save_params() + ["actor", "critic", "critic_target"]
-        Returns the names of the parameters that should be excluded by default
-        when saving the model.
 
-        :return: (List[str]) List of parameters that should be excluded from save
+    def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
-        """
-        # 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
-        """
         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")
-        return state_dicts, saved_tensors
+        return state_dicts, saved_pytorch_variables

stable_baselines3/td3/td3.py

@@ -205,19 +205,9 @@ class TD3(OffPolicyAlgorithm):
             reset_num_timesteps=reset_num_timesteps,
         )
 
-    def excluded_save_params(self) -> List[str]:
+    def _excluded_save_params(self) -> List[str]:
-        """
+        return super(TD3, self)._excluded_save_params() + ["actor", "critic", "actor_target", "critic_target"]
-        Returns the names of the parameters that should be excluded by default
-        when saving the model.
 
-        :return: (List[str]) List of parameters that should be excluded from save
+    def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
-        """
-        # 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
-        """
         state_dicts = ["policy", "actor.optimizer", "critic.optimizer"]
         return state_dicts, []

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
+    # Get parameters of different objects
-    params = deepcopy(model.policy.state_dict())
+    # deepcopy to avoid referencing to tensors we are about to modify
+    original_params = deepcopy(model.get_parameters())
 
-    # Modify all parameters to be random values
+    # Test different error cases of set_parameters.
-    random_params = dict((param_name, th.rand_like(param)) for param_name, param in params.items())
+    # 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()
+    new_params = model.get_parameters()
-    # Check that all params are different now
+    # Check that all params except the final item in each state-dict are different.
-    for k in params:
+    for object_name in original_params:
-        assert not th.allclose(params[k], new_params[k]), "Parameters did not change as expected."
+        # 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:
+        for object_name in new_params:
-            assert new_params[key].device.type == get_device(device).type
+            # Skip optimizers (no valid comparison with just th.allclose)
-            assert th.allclose(
+            if "optim" in object_name:
-                params[key].to("cpu"), new_params[key].to("cpu")
+                continue
-            ), "Model parameters not the same after save and load."
+            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)