dodox/metastable-baselines2
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Support for seting act_fn as string

Browse Source
This commit is contained in:
Dominik Moritz Roth 2023-09-05 19:12:26 +02:00
parent 3c81a15630
commit f0cd88365e
1 changed files with 15 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
sbBrix/common/policies.py
View File

@ -451,6 +451,11 @@ class ActorCriticPolicy(BasePolicy):
if optimizer_class == th.optim.Adam: if optimizer_class == th.optim.Adam:
optimizer_kwargs["eps"] = 1e-5 optimizer_kwargs["eps"] = 1e-5
if activation_fn == 'ReLU':
activation_fn = nn.ReLU
elif activation_fn == 'tanh':
activation_fn = nn.Tanh
super().__init__( super().__init__(
observation_space, observation_space,
action_space, action_space,
@ -791,7 +796,11 @@ class Actor(BasePolicy):
squash_output=True, squash_output=True,
) )
# Save arguments to re-create object at loading if activation_fn == 'ReLU':
activation_fn = nn.ReLU
elif activation_fn == 'tanh':
activation_fn = nn.Tanh
self.use_sde = use_sde self.use_sde = use_sde
self.use_pca = use_pca self.use_pca = use_pca
self.sde_features_extractor = None self.sde_features_extractor = None
@ -977,6 +986,11 @@ class SACPolicy(BasePolicy):
normalize_images=normalize_images, normalize_images=normalize_images,
) )
if activation_fn == 'ReLU':
activation_fn = nn.ReLU
elif activation_fn == 'tanh':
activation_fn = nn.Tanh
if net_arch is None: if net_arch is None:
net_arch = [256, 256] net_arch = [256, 256]