Merge branch 'develop-lin-enc-def' into develop-centralizedcritic-mm

Author: Ervin Teng

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

@@ -117,7 +117,13 @@ def __init__(
             self.self_size = 0
         self.ent_encoders = torch.nn.ModuleList(
             [
-                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
             ]
         )

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

@@ -120,14 +120,21 @@ class LinearEncoder(torch.nn.Module):
     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,
+    ):
         super().__init__()
         self.layers = [
             linear_layer(
                 input_size,
                 hidden_size,
-                kernel_init=Initialization.KaimingHeNormal,
-                kernel_gain=1.0,
+                kernel_init=kernel_init,
+                kernel_gain=kernel_gain,
             )
         ]
         self.layers.append(Swish())
@@ -136,8 +143,8 @@ def __init__(self, input_size: int, num_layers: int, hidden_size: int):
                 linear_layer(
                     hidden_size,
                     hidden_size,
-                    kernel_init=Initialization.KaimingHeNormal,
-                    kernel_gain=1.0,
+                    kernel_init=kernel_init,
+                    kernel_gain=kernel_gain,
                 )
             )
             self.layers.append(Swish())