From 1564a8508175ae152a6ba7bf51fa186030abae85 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 18 Oct 2021 10:43:56 +0200 Subject: [PATCH] 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 | 6 ++++ .github/ISSUE_TEMPLATE/custom_env.md | 6 ++++ docs/guide/examples.rst | 3 ++ docs/guide/save_format.rst | 10 ++++++- docs/misc/changelog.rst | 5 ++-- stable_baselines3/__init__.py | 1 + stable_baselines3/common/base_class.py | 12 +++++++- .../common/off_policy_algorithm.py | 2 +- stable_baselines3/common/save_util.py | 23 ++++++++++++-- stable_baselines3/common/utils.py | 30 ++++++++++++++++++- stable_baselines3/ddpg/ddpg.py | 2 +- stable_baselines3/dqn/dqn.py | 2 +- stable_baselines3/ppo/ppo.py | 2 -- stable_baselines3/sac/sac.py | 2 +- stable_baselines3/td3/td3.py | 2 +- stable_baselines3/version.txt | 2 +- tests/test_save_load.py | 6 +++- tests/test_utils.py | 13 +++++++- 18 files changed, 111 insertions(+), 18 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md index f67b4f2..7cbc198 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.md +++ b/.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. diff --git a/.github/ISSUE_TEMPLATE/custom_env.md b/.github/ISSUE_TEMPLATE/custom_env.md index bea9b59..0a12a68 100644 --- a/.github/ISSUE_TEMPLATE/custom_env.md +++ b/.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. diff --git a/docs/guide/examples.rst b/docs/guide/examples.rst index 7e870f6..929e587 100644 --- a/docs/guide/examples.rst +++ b/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 diff --git a/docs/guide/save_format.rst b/docs/guide/save_format.rst index 38dc233..8bd9aa8 100644 --- a/docs/guide/save_format.rst +++ b/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 `_ or `#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: diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 904e10e..8ff6173 100644 --- a/docs/misc/changelog.rst +++ b/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: ^^^^^^^^^^ diff --git a/stable_baselines3/__init__.py b/stable_baselines3/__init__.py index acca18a..4e31c5b 100644 --- a/stable_baselines3/__init__.py +++ b/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 diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 23f14e3..1b38555 100644 --- a/stable_baselines3/common/base_class.py +++ b/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: diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index c26c0a4..999365a 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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, diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 4b8c976..dcacfba 100644 --- a/stable_baselines3/common/save_util.py +++ b/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). diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index be5d488..7548cdc 100644 --- a/stable_baselines3/common/utils.py +++ b/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 diff --git a/stable_baselines3/ddpg/ddpg.py b/stable_baselines3/ddpg/ddpg.py index 01cdcda..14293ca 100644 --- a/stable_baselines3/ddpg/ddpg.py +++ b/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, diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 69b2227..d3b1eb5 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/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, diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index ab1129a..9e16e04 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/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) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 605f086..e9502fd 100644 --- a/stable_baselines3/sac/sac.py +++ b/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, diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 1eb28f7..e059761 100644 --- a/stable_baselines3/td3/td3.py +++ b/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, diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 90ebae4..16156e3 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -1.2.1a3 +1.2.1a4 diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 7b1fef5..1454b98 100644 --- a/tests/test_save_load.py +++ b/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"] diff --git a/tests/test_utils.py b/tests/test_utils.py index f4092d7..711176c 100644 --- a/tests/test_utils.py +++ b/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