add default args to LinearEncoder

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

                 # LinearEncoder(self.self_size + ent_size, 2, embedding_size)
                 # from http://www.cs.toronto.edu/~mvolkovs/ICML2020_tfixup.pdf
                 # linear_layer(self.self_size + ent_size, embedding_size, Initialization.Normal, kernel_gain=(.125 / (self.self_size + ent_size)) ** 0.5)
-                 linear_layer(ent_size, embedding_size, Initialization.Normal, kernel_gain=(.125 / (self.self_size + ent_size)) ** 0.5)
-        #         LinearEncoder(self.self_size + ent_size, 1, embedding_size, layer_norm=False)
+                 #linear_layer(ent_size, embedding_size, Initialization.Normal, kernel_gain=(.125 / (self.self_size + ent_size)) ** 0.5)
+                 LinearEncoder(ent_size, 1, embedding_size, kernel_init=Initialization.Normal, kernel_gain=(.125 / embedding_size) ** 0.5)
+                 #LinearEncoder(self.self_size + ent_size, 1, embedding_size, layer_norm=False)
                 for ent_size in self.entities_sizes
             ]
         )

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

     Linear layers.
     """
 
-    def __init__(self, input_size: int, num_layers: int, hidden_size: int, layer_norm=False):
+    def __init__(self, input_size: int, num_layers: int, hidden_size: int, kernel_init=Initialization.KaimingHeNormal, kernel_gain=1.0):
-                kernel_init=Initialization.KaimingHeNormal,
-                kernel_gain=1.0,
+                kernel_init=kernel_init,
+                kernel_gain=kernel_gain,
-        if layer_norm:
-            self.layers.append(torch.nn.LayerNorm(hidden_size, elementwise_affine=True))
-                    kernel_init=Initialization.KaimingHeNormal,
-                    kernel_gain=1.0,
+                    kernel_init=kernel_init,
+                    kernel_gain=kernel_gain,
-            if layer_norm:
-                self.layers.append(torch.nn.LayerNorm(hidden_size, elementwise_affine=True))
-
         self.seq_layers = torch.nn.Sequential(*self.layers)
 
     def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:

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

             tens.retain_grad()
         total_enc_size += encoded_act_size
         self.linear_encoder = LinearEncoder(
-            total_enc_size, n_layers, self.h_size, layer_norm=False
+            total_enc_size, n_layers, self.h_size
         )
         for _,tens in list(self.linear_encoder.named_parameters()):
             tens.retain_grad()