Fix error when loading a model that has net_arch manually set to None (#1937)

* Fix loading a model with net_arch=None

* Remove redundant get

* Dummy commit

* Add to contributors

* Update test and version

---------

Co-authored-by: Antonin Raffin <antonin.raffin@ensta.org>
This commit is contained in:
Joe Ksiazek 2024-06-05 11:27:40 -04:00 committed by GitHub
parent 6c00565778
commit 0b06d8ab20
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 19 additions and 7 deletions

View file

@ -3,7 +3,7 @@
Changelog
==========
Release 2.4.0a1 (WIP)
Release 2.4.0a2 (WIP)
--------------------------
Breaking Changes:
@ -17,6 +17,7 @@ Bug Fixes:
- Fixed memory leak when loading learner from storage, ``set_parameters()`` does not try to load the object data anymore
and only loads the PyTorch parameters (@peteole)
- Cast type in compute gae method to avoid error when using torch compile (@amjames)
- Fixed error when loading a model that has ``net_arch`` manually set to ``None`` (@jak3122)
`SB3-Contrib`_
^^^^^^^^^^^^^^
@ -1661,4 +1662,4 @@ And all the contributors:
@anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong @ReHoss
@DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError @harveybellini @DavyMorgan @FieteO @jonasreiher @npit @WeberSamuel @troiganto
@lutogniew @lbergmann1 @lukashass @BertrandDecoster @pseudo-rnd-thoughts @stefanbschneider @kyle-he @PatrickHelm @corentinlger
@marekm4 @stagoverflow @rushitnshah @markscsmith @NickLucche @cschindlbeck @peteole
@marekm4 @stagoverflow @rushitnshah @markscsmith @NickLucche @cschindlbeck @peteole @jak3122

View file

@ -692,10 +692,9 @@ class BaseAlgorithm(ABC):
if "device" in data["policy_kwargs"]:
del data["policy_kwargs"]["device"]
# backward compatibility, convert to new format
if "net_arch" in data["policy_kwargs"] and len(data["policy_kwargs"]["net_arch"]) > 0:
saved_net_arch = data["policy_kwargs"]["net_arch"]
if isinstance(saved_net_arch, list) and isinstance(saved_net_arch[0], dict):
data["policy_kwargs"]["net_arch"] = saved_net_arch[0]
saved_net_arch = data["policy_kwargs"].get("net_arch")
if saved_net_arch and isinstance(saved_net_arch, list) and isinstance(saved_net_arch[0], dict):
data["policy_kwargs"]["net_arch"] = saved_net_arch[0]
if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]:
raise ValueError(

View file

@ -1 +1 @@
2.4.0a1
2.4.0a2

View file

@ -797,3 +797,15 @@ def test_cast_lr_schedule(tmp_path):
model = PPO.load(tmp_path / "ppo.zip")
assert type(model.lr_schedule(1.0)) is float # noqa: E721
assert np.allclose(model.lr_schedule(0.5), 0.5 * np.sin(1.0))
def test_save_load_net_arch_none(tmp_path):
"""
Test that the model is loaded correctly when net_arch is manually set to None.
See GH#1928
"""
PPO("MlpPolicy", "CartPole-v1", policy_kwargs=dict(net_arch=None)).save(tmp_path / "ppo.zip")
model = PPO.load(tmp_path / "ppo.zip")
# None has been replaced by the default net arch
assert model.policy.net_arch is not None
os.remove(tmp_path / "ppo.zip")