Parameterize hidden layer size of policy, add some candidates to grid search

2022-02-23 18:14:23 +01:00
parent 7f64ec7bb0
commit 336cf02278
2 changed files with 8 additions and 7 deletions
--- a/src/core/policy.py
+++ b/src/core/policy.py
@@ -51,12 +51,12 @@ class Policy(BasePolicy):

 class DiscretePolicy(BasePolicy):

-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, hidden_layer_size=50, **kwargs):
        super().__init__(*args, **kwargs)
        self.nn = nn.Sequential(
-            nn.LazyLinear(50),
+            nn.LazyLinear(hidden_layer_size),
            nn.Tanh(),
-            nn.LazyLinear(50),
+            nn.LazyLinear(hidden_layer_size),
            nn.Tanh(),
            nn.LazyLinear(self.action_dim),
        )