From 723b341c61d168e1460399592d5cebd4c6ef3cc8 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Tue, 24 Nov 2020 15:13:00 +0100 Subject: [PATCH] Fix for saving big replay buffer, use pickle protocol>=4 (#239) --- docs/misc/changelog.rst | 3 ++- stable_baselines3/common/save_util.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 014630e..b154e49 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -32,6 +32,7 @@ Bug Fixes: - Fixed ``DQN`` predict method when using single ``gym.Env`` with ``deterministic=False`` - Fixed bug that the arguments order of ``explained_variance()`` in ``ppo.py`` and ``a2c.py`` is not correct (@thisray) - Fixed bug where full ``HerReplayBuffer`` leads to an index error. (@megan-klaiber) +- Fixed bug where replay buffer could not be saved if it was too big (> 4 Gb) for python<3.8 (thanks @hn2) Deprecations: ^^^^^^^^^^^^^ @@ -523,4 +524,4 @@ And all the contributors: @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio @diditforlulz273 @liorcohen5 @ManifoldFR @mloo3 @SwamyDev @wmmc88 @megan-klaiber @thisray -@tfederico +@tfederico @hn2 diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 4d9ed8b..4b2f1d6 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -337,7 +337,9 @@ def save_to_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], obj: Any, ver :param verbose: Verbosity level, 0 means only warnings, 2 means debug information. """ with open_path(path, "w", verbose=verbose, suffix="pkl") as file_handler: - pickle.dump(obj, file_handler) + # Use protocol>=4 to support saving replay buffers >= 4Gb + # See https://docs.python.org/3/library/pickle.html + pickle.dump(obj, file_handler, protocol=pickle.HIGHEST_PROTOCOL) def load_from_pkl(path: Union[str, pathlib.Path, io.BufferedIOBase], verbose: int = 0) -> Any: