fix: squash the mean before gaussian for stability (tinkoff-ai#41)

* fix: squash the mean before gaussian for stability * make the linter happy * make the linter happy * delete whitespace
jiawei415 · Mar 29, 2023 · b62fa28 · b62fa28
1 parent 2a7b88c
commit b62fa28
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/algorithms/iql.py b/algorithms/iql.py
@@ -279,7 +279,10 @@ def __init__(
         n_hidden: int = 2,
     ):
         super().__init__()
-        self.net = MLP([state_dim, *([hidden_dim] * n_hidden), act_dim])
+        self.net = MLP(
+            [state_dim, *([hidden_dim] * n_hidden), act_dim],
+            output_activation_fn=nn.Tanh,
+        )
         self.log_std = nn.Parameter(torch.zeros(act_dim, dtype=torch.float32))
         self.max_action = max_action