saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-27 20:02:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix for new version

Browse source
This commit is contained in:
Antonin Raffin 2020-09-24 15:25:03 +02:00
parent 4112e27a1e
commit 243f09837c
2 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
stable_baselines3/common/save_util.py
View file

@ -404,7 +404,7 @@ def load_from_zip_file(
# Load the parameters with the right ``map_location``.
# Remove ".pth" ending with splitext
th_object = th.load(file_content, map_location=device)
if file_path == "pytorch_variables.pth":
if file_path == "pytorch_variables.pth" or file_path == "tensors.pth":
# PyTorch variables (not state_dicts)
pytorch_variables = th_object
else:

10
stable_baselines3/tqc/tqc.py
View file

@ -468,7 +468,7 @@ class TQC(OffPolicyAlgorithm):
reset_num_timesteps=reset_num_timesteps,
)
def excluded_save_params(self) -> List[str]:
def _excluded_save_params(self) -> List[str]:
"""
Returns the names of the parameters that should be excluded by default
when saving the model.
@ -478,14 +478,14 @@ class TQC(OffPolicyAlgorithm):
# Exclude aliases
return super(TQC, self).excluded_save_params() + ["actor", "critic", "critic_target"]
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", "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")
return state_dicts, saved_tensors
saved_pytorch_variables.append("ent_coef_tensor")
return state_dicts, saved_pytorch_variables