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

Fix custom optimizer

Browse source
This commit is contained in:
Antonin RAFFIN 2020-04-22 13:21:11 +02:00
parent 041f2bc59a
commit f3cb0688c4
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/test_custom_policy.py
View file

@ -30,5 +30,5 @@ def test_custom_offpolicy(model_class, net_arch):
@pytest.mark.parametrize('model_class', [A2C, PPO, SAC, TD3])
@pytest.mark.parametrize('optimizer_kwargs', [None, dict(weight_decay=0.0)])
def test_custom_optimizer(model_class, optimizer_kwargs):
policy_kwargs = dict(optimizer=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32])
policy_kwargs = dict(optimizer_class=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32])
_ = model_class('MlpPolicy', 'Pendulum-v0', policy_kwargs=policy_kwargs).learn(1000)