huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added support for different activations, number of hidden layers, options, running a big run over this

Browse Source
This commit is contained in:
Arec Jamgochian
2022-02-25 00:00:04 -08:00
parent 7feea74eb8
commit 99aa50a08b
3 changed files with 27 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/core/policy.py
View File

@@ -51,15 +51,18 @@ class Policy(BasePolicy):
class DiscretePolicy(BasePolicy):
def __init__(self, *args, hidden_layer_size=50, **kwargs):
def __init__(self, *args, hidden_layer_size=50, n_hidden_layers=2, activation=nn.Tanh, **kwargs):
super().__init__(*args, **kwargs)
self.nn = nn.Sequential(
nn.LazyLinear(hidden_layer_size),
nn.Tanh(),
nn.LazyLinear(hidden_layer_size),
nn.Tanh(),
nn.LazyLinear(self.action_dim),
)
layers = [nn.LazyLinear(hidden_layer_size), activation()] * n_hidden_layers
self.nn = nn.Sequential(*layers, nn.LazyLinear(self.action_dim))
#self.nn = nn.Sequential(
# nn.LazyLinear(hidden_layer_size),
# nn.Tanh(),
# nn.LazyLinear(hidden_layer_size),
# nn.Tanh(),
# nn.LazyLinear(self.action_dim),
#)
def forward(self, states):
return self.nn(states)