System info helper (#613)

* Add `system_env_info`

* Add `print_system_info` to load
and store system info at save time

* Remove TODO

* Rename to `get_system_info`

* Import as sb3 for consistency

* Update changelog

* Add warning for old SB3 versions

* Use underscore litteral for more clarity

.github/ISSUE_TEMPLATE/bug_report.md

@@ -50,6 +50,12 @@ Describe the characteristic of your environment:
  * Gym version
  * Versions of any other relevant libraries
 
+You can use `sb3.get_system_info()` to print relevant packages info:
+```python
+import stable_baselines3 as sb3
+sb3.get_system_info()
+```
+
 ### Additional context
 Add any other context about the problem here.
 

.github/ISSUE_TEMPLATE/custom_env.md

@@ -78,6 +78,12 @@ Describe the characteristic of your environment:
  * Gym version
  * Versions of any other relevant libraries
 
+You can use `sb3.get_system_info()` to print relevant packages info:
+```python
+import stable_baselines3 as sb3
+sb3.get_system_info()
+```
+
 ### Additional context
 Add any other context about the problem here.
 

docs/guide/examples.rst

@@ -81,6 +81,9 @@ In the following example, we will train, save and load a DQN model on the Lunar
   del model  # delete trained model to demonstrate loading
 
   # Load the trained agent
+  # NOTE: if you have loading issue, you can pass `print_system_info=True`
+  # to compare the system on which the model was trained vs the current one
+  # model = DQN.load("dqn_lunar", env=env, print_system_info=True)
   model = DQN.load("dqn_lunar", env=env)
 
   # Evaluate the agent

docs/guide/save_format.rst

@@ -30,8 +30,15 @@ inspecting stored objects without deserializing the object itself.
 
 This format allows skipping elements in the file, i.e. we can skip deserializing objects that are
 broken/non-serializable.
+This can be done via ``custom_objects`` argument to load functions.
 
-.. This can be done via ``custom_objects`` argument to load functions.
+.. note::
+
+  If you encounter loading issue, for instance pickle issues or error after loading
+  (see `#171 <https://github.com/DLR-RM/stable-baselines3/issues/171>`_ or `#573 <https://github.com/DLR-RM/stable-baselines3/issues/573>`_),
+  you can pass ``print_system_info=True``
+  to compare the system on which the model was trained vs the current one
+  ``model = PPO.load("ppo_saved", print_system_info=True)``
 
 
 File structure:
@@ -44,6 +51,7 @@ File structure:
   ├── policy.pth        PyTorch state dictionary of the policy saved
   ├── pytorch_variables.pth Additional PyTorch variables
   ├── _stable_baselines3_version contains the SB3 version with which the model was saved
+  ├── system_info.txt contains system info (os, python version, ...) on which the model was saved
 
 
 Pros:

docs/misc/changelog.rst

@@ -4,10 +4,9 @@ Changelog
 ==========
 
 
-Release 1.2.1a3 (WIP)
+Release 1.2.1a4 (WIP)
 ---------------------------
 
-
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - ``sde_net_arch`` argument in policies is deprecated and will be removed in a future version.
@@ -22,6 +21,8 @@ New Features:
 ^^^^^^^^^^^^^
 - Added methods ``get_distribution`` and ``predict_values`` for ``ActorCriticPolicy`` for A2C/PPO/TRPO (@cyprienc)
 - Added methods ``forward_actor`` and ``forward_critic`` for ``MlpExtractor``
+- Added ``sb3.get_system_info()`` helper function to gather version information relevant to SB3 (e.g., Python and PyTorch version)
+- Saved models now store system information where agent was trained, and load functions have ``print_system_info`` parameter to help debugging load issues.
 
 Bug Fixes:
 ^^^^^^^^^^

stable_baselines3/__init__.py

@@ -1,6 +1,7 @@
 import os
 
 from stable_baselines3.a2c import A2C
+from stable_baselines3.common.utils import get_system_info
 from stable_baselines3.ddpg import DDPG
 from stable_baselines3.dqn import DQN
 from stable_baselines3.her.her_replay_buffer import HerReplayBuffer

stable_baselines3/common/base_class.py

@@ -25,6 +25,7 @@ from stable_baselines3.common.utils import (
     check_for_correct_spaces,
     get_device,
     get_schedule_fn,
+    get_system_info,
     set_random_seed,
     update_learning_rate,
 )
@@ -634,6 +635,7 @@ class BaseAlgorithm(ABC):
         env: Optional[GymEnv] = None,
         device: Union[th.device, str] = "auto",
         custom_objects: Optional[Dict[str, Any]] = None,
+        print_system_info: bool = False,
         **kwargs,
     ) -> "BaseAlgorithm":
         """
@@ -650,9 +652,17 @@ class BaseAlgorithm(ABC):
             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.
+        :param print_system_info: Whether to print system info from the saved model
+            and the current system info (useful to debug loading issues)
         :param kwargs: extra arguments to change the model when loading
         """
-        data, params, pytorch_variables = load_from_zip_file(path, device=device, custom_objects=custom_objects)
+        if print_system_info:
+            print("== CURRENT SYSTEM INFO ==")
+            get_system_info()
+
+        data, params, pytorch_variables = load_from_zip_file(
+            path, device=device, custom_objects=custom_objects, print_system_info=print_system_info
+        )
 
         # Remove stored device information and replace with ours
         if "policy_kwargs" in data:

stable_baselines3/common/off_policy_algorithm.py

@@ -79,7 +79,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
         env: Union[GymEnv, str],
         policy_base: Type[BasePolicy],
         learning_rate: Union[float, Schedule],
-        buffer_size: int = 1000000,  # 1e6
+        buffer_size: int = 1_000_000,  # 1e6
         learning_starts: int = 100,
         batch_size: int = 256,
         tau: float = 0.005,

stable_baselines3/common/save_util.py

@@ -16,9 +16,9 @@ from typing import Any, Dict, Optional, Tuple, Union
 import cloudpickle
 import torch as th
 
-import stable_baselines3
+import stable_baselines3 as sb3
 from stable_baselines3.common.type_aliases import TensorDict
-from stable_baselines3.common.utils import get_device
+from stable_baselines3.common.utils import get_device, get_system_info
 
 
 def recursive_getattr(obj: Any, attr: str, *args) -> Any:
@@ -321,7 +321,9 @@ def save_to_zip_file(
                 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__)
+        archive.writestr("_stable_baselines3_version", sb3.__version__)
+        # Save system info about the current python env
+        archive.writestr("system_info.txt", get_system_info(print_info=False)[1])
 
 
 def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj: Any, verbose: int = 0) -> None:
@@ -362,6 +364,7 @@ def load_from_zip_file(
     custom_objects: Optional[Dict[str, Any]] = None,
     device: Union[th.device, str] = "auto",
     verbose: int = 0,
+    print_system_info: bool = False,
 ) -> (Tuple[Optional[Dict[str, Any]], Optional[TensorDict], Optional[TensorDict]]):
     """
     Load model data from a .zip archive
@@ -376,6 +379,9 @@ def load_from_zip_file(
         ``keras.models.load_model``. Useful when you have an object in
         file that can not be deserialized.
     :param device: Device on which the code should run.
+    :param verbose: Verbosity level, 0 means only warnings, 2 means debug information.
+    :param print_system_info: Whether to print or not the system info
+        about the saved model.
     :return: Class parameters, model state_dicts (aka "params", dict of state_dict)
         and dict of pytorch variables
     """
@@ -395,6 +401,17 @@ def load_from_zip_file(
             pytorch_variables = None
             params = {}
 
+            # Debug system info first
+            if print_system_info:
+                if "system_info.txt" in namelist:
+                    print("== SAVED MODEL SYSTEM INFO ==")
+                    print(archive.read("system_info.txt").decode())
+                else:
+                    warnings.warn(
+                        "The model was saved with SB3 <= 1.2.0 and thus cannot print system information.",
+                        UserWarning,
+                    )
+
             if "data" in namelist and load_data:
                 # Load class parameters that are stored
                 # with either JSON or pickle (not PyTorch variables).

stable_baselines3/common/utils.py

@@ -1,14 +1,17 @@
 import glob
 import os
+import platform
 import random
 from collections import deque
 from itertools import zip_longest
-from typing import Dict, Iterable, Optional, Union
+from typing import Dict, Iterable, Optional, Tuple, Union
 
 import gym
 import numpy as np
 import torch as th
 
+import stable_baselines3 as sb3
+
 # Check if tensorboard is available for pytorch
 try:
     from torch.utils.tensorboard import SummaryWriter
@@ -460,3 +463,28 @@ def should_collect_more_steps(
             "The unit of the `train_freq` must be either TrainFrequencyUnit.STEP "
             f"or TrainFrequencyUnit.EPISODE not '{train_freq.unit}'!"
         )
+
+
+def get_system_info(print_info: bool = True) -> Tuple[Dict[str, str], str]:
+    """
+    Retrieve system and python env info for the current system.
+
+    :param print_info: Whether to print or not those infos
+    :return: Dictionary summing up the version for each relevant package
+        and a formatted string.
+    """
+    env_info = {
+        "OS": f"{platform.platform()} {platform.version()}",
+        "Python": platform.python_version(),
+        "Stable-Baselines3": sb3.__version__,
+        "PyTorch": th.__version__,
+        "GPU Enabled": str(th.cuda.is_available()),
+        "Numpy": np.__version__,
+        "Gym": gym.__version__,
+    }
+    env_info_str = ""
+    for key, value in env_info.items():
+        env_info_str += f"{key}: {value}\n"
+    if print_info:
+        print(env_info_str)
+    return env_info, env_info_str

stable_baselines3/ddpg/ddpg.py

@@ -58,7 +58,7 @@ class DDPG(TD3):
         policy: Union[str, Type[TD3Policy]],
         env: Union[GymEnv, str],
         learning_rate: Union[float, Schedule] = 1e-3,
-        buffer_size: int = 1000000,  # 1e6
+        buffer_size: int = 1_000_000,  # 1e6
         learning_starts: int = 100,
         batch_size: int = 100,
         tau: float = 0.005,

stable_baselines3/dqn/dqn.py

@@ -63,7 +63,7 @@ class DQN(OffPolicyAlgorithm):
         policy: Union[str, Type[DQNPolicy]],
         env: Union[GymEnv, str],
         learning_rate: Union[float, Schedule] = 1e-4,
-        buffer_size: int = 1000000,  # 1e6
+        buffer_size: int = 1_000_000,  # 1e6
         learning_starts: int = 50000,
         batch_size: Optional[int] = 32,
         tau: float = 1.0,

stable_baselines3/ppo/ppo.py

@@ -193,8 +193,6 @@ class PPO(OnPolicyAlgorithm):
                     actions = rollout_data.actions.long().flatten()
 
                 # Re-sample the noise matrix because the log_std has changed
-                # TODO: investigate why there is no issue with the gradient
-                # if that line is commented (as in SAC)
                 if self.use_sde:
                     self.policy.reset_noise(self.batch_size)
 

stable_baselines3/sac/sac.py

@@ -77,7 +77,7 @@ class SAC(OffPolicyAlgorithm):
         policy: Union[str, Type[SACPolicy]],
         env: Union[GymEnv, str],
         learning_rate: Union[float, Schedule] = 3e-4,
-        buffer_size: int = 1000000,  # 1e6
+        buffer_size: int = 1_000_000,  # 1e6
         learning_starts: int = 100,
         batch_size: int = 256,
         tau: float = 0.005,

stable_baselines3/td3/td3.py

@@ -65,7 +65,7 @@ class TD3(OffPolicyAlgorithm):
         policy: Union[str, Type[TD3Policy]],
         env: Union[GymEnv, str],
         learning_rate: Union[float, Schedule] = 1e-3,
-        buffer_size: int = 1000000,  # 1e6
+        buffer_size: int = 1_000_000,  # 1e6
         learning_starts: int = 100,
         batch_size: int = 100,
         tau: float = 0.005,

stable_baselines3/version.txt

@@ -1 +1 @@
-1.2.1a3
+1.2.1a4

tests/test_save_load.py

@@ -222,7 +222,11 @@ def test_exclude_include_saved_params(tmp_path, model_class):
     del model
     # Load with custom objects
     custom_objects = dict(learning_rate=2e-5, dummy=1.0)
-    model = model_class.load(str(tmp_path / "test_save.zip"), custom_objects=custom_objects)
+    model = model_class.load(
+        str(tmp_path / "test_save.zip"),
+        custom_objects=custom_objects,
+        print_system_info=True,
+    )
     assert model.verbose == 2
     # Check that the custom object was taken into account
     assert model.learning_rate == custom_objects["learning_rate"]

tests/test_utils.py

@@ -6,13 +6,14 @@ import numpy as np
 import pytest
 import torch as th
 
+import stable_baselines3 as sb3
 from stable_baselines3 import A2C, PPO
 from stable_baselines3.common.atari_wrappers import ClipRewardEnv, MaxAndSkipEnv
 from stable_baselines3.common.env_util import is_wrapped, make_atari_env, make_vec_env, unwrap_wrapper
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.monitor import Monitor
 from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise
-from stable_baselines3.common.utils import polyak_update, zip_strict
+from stable_baselines3.common.utils import get_system_info, polyak_update, zip_strict
 from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv
 
 
@@ -376,3 +377,13 @@ def test_ppo_warnings():
     # Truncated mini-batch
     with pytest.warns(UserWarning):
         PPO("MlpPolicy", "Pendulum-v0", n_steps=6, batch_size=8)
+
+
+def test_get_system_info():
+    info, info_str = get_system_info(print_info=True)
+    assert info["Stable-Baselines3"] == str(sb3.__version__)
+    assert "Python" in info_str
+    assert "PyTorch" in info_str
+    assert "GPU Enabled" in info_str
+    assert "Numpy" in info_str
+    assert "Gym" in info_str