add LayerNorm (#4847)

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

     lstm_layer,
     Initialization,
     LSTM,
+    LayerNorm,
 )
 
 
     # Hidden size should be half of memory_size
     assert out.shape == (batch_size, seq_len, memory_size // 2)
     assert mem.shape == (1, batch_size, memory_size)
+
+
+def test_layer_norm():
+    torch.manual_seed(0)
+    torch_ln = torch.nn.LayerNorm(10, elementwise_affine=False)
+    cust_ln = LayerNorm()
+
+    sample_input = torch.rand(10)
+    assert torch.all(
+        torch.isclose(
+            torch_ln(sample_input), cust_ln(sample_input), atol=1e-5, rtol=0.0
+        )
+    )
+    sample_input = torch.rand((4, 10))
+    assert torch.all(
+        torch.isclose(
+            torch_ln(sample_input), cust_ln(sample_input), atol=1e-5, rtol=0.0
+        )
+    )
+    sample_input = torch.rand((7, 6, 10))
+    assert torch.all(
+        torch.isclose(
+            torch_ln(sample_input), cust_ln(sample_input), atol=1e-5, rtol=0.0
+        )
+    )

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

 from mlagents.torch_utils import torch
 from typing import Tuple, Optional, List
-from mlagents.trainers.torch.layers import LinearEncoder, Initialization, linear_layer
+from mlagents.trainers.torch.layers import (
+    LinearEncoder,
+    Initialization,
+    linear_layer,
+    LayerNorm,
+)
 from mlagents.trainers.torch.model_serialization import exporting_to_onnx
 from mlagents.trainers.exception import UnityTrainerException
 
                 for ent_size in self.entity_sizes
             ]
         )
+        self.embedding_norm = LayerNorm()
 
     def forward(
         self, x_self: torch.Tensor, entities: List[torch.Tensor]
             [ent_encoder(x) for ent_encoder, x in zip(self.ent_encoders, self_and_ent)],
             dim=1,
         )
+        encoded_entities = self.embedding_norm(encoded_entities)
         return encoded_entities
 
     @staticmethod
             kernel_init=Initialization.Normal,
             kernel_gain=(0.125 / embedding_size) ** 0.5,
         )
+        self.residual_norm = LayerNorm()
 
     def forward(self, inp: torch.Tensor, key_masks: List[torch.Tensor]) -> torch.Tensor:
         # Gather the maximum number of entities information
         output, _ = self.attention(query, key, value, num_ent, num_ent, mask)
         # Residual
         output = self.fc_out(output) + inp
+        output = self.residual_norm(output)
-        # Residual between x_self and the output of the module
         return output

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

         pass
 
 
+class LayerNorm(torch.nn.Module):
+    """
+    A vanilla implementation of layer normalization  https://arxiv.org/pdf/1607.06450.pdf
+    norm_x = (x - mean) / sqrt((x - mean) ^ 2)
+    This does not include the trainable parameters gamma and beta for performance speed.
+    Typically, this is norm_x * gamma + beta
+    """
+
+    def forward(self, layer_activations: torch.Tensor) -> torch.Tensor:
+        mean = torch.mean(layer_activations, dim=-1, keepdim=True)
+        var = torch.mean((layer_activations - mean) ** 2, dim=-1, keepdim=True)
+        return (layer_activations - mean) / (torch.sqrt(var + 1e-5))
+
+
 class LinearEncoder(torch.nn.Module):
     """
     Linear layers.