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

fixing issues with lazylinear sequential, setting off a big run

Browse Source
This commit is contained in:
Arec Jamgochian
2022-02-25 17:48:56 -08:00
parent 2dfd7e3c2b
commit fa98601fa6
2 changed files with 19 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/core/policy.py
View File

@@ -36,15 +36,20 @@ class BasePolicy(nn.Module):
class Policy(BasePolicy):
def __init__(self, *args, **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(50),
nn.Tanh(),
nn.LazyLinear(50),
nn.Tanh(),
nn.LazyLinear(2 * self.action_dim),
)
layers = sum([[nn.LazyLinear(hidden_layer_size),
activation()] for _ in range(n_hidden_layers)],[])
self.nn = nn.Sequential(*layers, nn.LazyLinear(2 *self.action_dim))
# old
# self.nn = nn.Sequential(
# nn.LazyLinear(50),
# nn.Tanh(),
# nn.LazyLinear(50),
# nn.Tanh(),
# nn.LazyLinear(2 * self.action_dim),
#)
def forward(self, states):
return self.nn(states)
@@ -53,17 +58,10 @@ class DiscretePolicy(BasePolicy):
def __init__(self, *args, hidden_layer_size=50, n_hidden_layers=2, activation=nn.Tanh, **kwargs):
super().__init__(*args, **kwargs)
layers = [nn.LazyLinear(hidden_layer_size), activation()] * n_hidden_layers
layers = sum([[nn.LazyLinear(hidden_layer_size),
activation()] for _ in range(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)