diff --git a/network/network_bodies.py b/network/network_bodies.py index ef60fc9..ad73360 100644 --- a/network/network_bodies.py +++ b/network/network_bodies.py @@ -24,7 +24,7 @@ class NatureConvBody(nn.Module): return y class FCBody(nn.Module): - def __init__(self, state_dim, hidden_units=(64, 64), gate=F.tanh): + def __init__(self, state_dim, hidden_units=(64, 64), gate=F.relu): super(FCBody, self).__init__() dims = (state_dim, ) + hidden_units self.layers = nn.ModuleList([layer_init(nn.Linear(dim_in, dim_out)) for dim_in, dim_out in zip(dims[:-1], dims[1:])])