Add default init/gain to LinearEncoder (#4846)

ml-agents/mlagents/trainers/tests/torch/test_attention.py

         list(transformer.parameters()) + list(l_layer.parameters()), lr=0.001
     )
     batch_size = 200
-    point_range = 3
-    init_error = -1.0
-    for _ in range(250):
-        center = torch.rand((batch_size, size)) * point_range * 2 - point_range
-        key = torch.rand((batch_size, n_k, size)) * point_range * 2 - point_range
+    for _ in range(200):
+        center = torch.rand((batch_size, size))
+        key = torch.rand((batch_size, n_k, size))
         with torch.no_grad():
             # create the target : The key closest to the query in euclidean distance
             distance = torch.sum(
         prediction = prediction.reshape((batch_size, size))
         error = torch.mean((prediction - target) ** 2, dim=1)
         error = torch.mean(error) / 2
-        if init_error == -1.0:
-            init_error = error.item()
-        else:
-            assert error.item() < init_error
-    assert error.item() < 0.3
+    assert error.item() < 0.02

ml-agents/mlagents/trainers/torch/attention.py

         # If not concatenating self, input to encoder is just entity size
         if not concat_self:
             self.self_size = 0
+        # Initialization scheme from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf
-                LinearEncoder(self.self_size + ent_size, 1, embedding_size)
+                LinearEncoder(
+                    self.self_size + ent_size,
+                    1,
+                    embedding_size,
+                    kernel_init=Initialization.Normal,
+                    kernel_gain=(0.125 / embedding_size) ** 0.5,
+                )
                 for ent_size in self.entity_sizes
             ]
         )
             num_heads=num_heads, embedding_size=embedding_size
         )
 
+        # Initialization scheme from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf
         self.fc_q = linear_layer(
             embedding_size,
             embedding_size,

ml-agents/mlagents/trainers/torch/layers.py

     Linear layers.
     """
 
-    def __init__(self, input_size: int, num_layers: int, hidden_size: int):
+    def __init__(
+        self,
+        input_size: int,
+        num_layers: int,
+        hidden_size: int,
+        kernel_init: Initialization = Initialization.KaimingHeNormal,
+        kernel_gain: float = 1.0,
+    ):
-                kernel_init=Initialization.KaimingHeNormal,
-                kernel_gain=1.0,
+                kernel_init=kernel_init,
+                kernel_gain=kernel_gain,
             )
         ]
         self.layers.append(Swish())
                     hidden_size,
                     hidden_size,
-                    kernel_init=Initialization.KaimingHeNormal,
-                    kernel_gain=1.0,
+                    kernel_init=kernel_init,
+                    kernel_gain=kernel_gain,
                 )
             )
             self.layers.append(Swish())