Experimental amrl layer

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

             _init_methods[bias_init](param.data)
             param.data[hidden_size : 2 * hidden_size].add_(forget_bias)
     return lstm
+
+
+class AMRLMax(torch.nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        num_layers: int = 1,
+        batch_first: bool = True,
+        forget_bias: float = 1.0,
+        kernel_init: Initialization = Initialization.XavierGlorotUniform,
+        bias_init: Initialization = Initialization.Zero,
+    ):
+        super().__init__()
+        self.lstm = lstm_layer(
+            input_size,
+            hidden_size,
+            num_layers,
+            batch_first,
+            forget_bias,
+            kernel_init,
+            bias_init,
+        )
+        self.hidden_size = hidden_size
+
+    def forward(self, input_tensor, h0_c0):
+        hidden = h0_c0
+        all_out = []
+        m = None
+        for t in range(input_tensor.shape[1]):
+            out, hidden = self.lstm(input_tensor[:, t : t + 1, :], hidden)
+            h_half, other_half = torch.split(out, self.hidden_size // 2, dim=-1)
+            if m is None:
+                m = h_half
+            else:
+                m = torch.max(m, h_half)
+            out = torch.cat([m, other_half])
+            all_out.append(out)
+        return torch.cat(all_out, dim=1), hidden

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

 from mlagents.trainers.settings import NetworkSettings
 from mlagents.trainers.torch.utils import ModelUtils
 from mlagents.trainers.torch.decoders import ValueHeads
-from mlagents.trainers.torch.layers import lstm_layer
+from mlagents.trainers.torch.layers import AMRLMax
 
 ActivationFunction = Callable[[torch.Tensor], torch.Tensor]
 EncoderFunction = Callable[
         )
 
         if self.use_lstm:
-            self.lstm = lstm_layer(self.h_size, self.m_size // 2, batch_first=True)
+            self.lstm = AMRLMax(self.h_size, self.m_size // 2, batch_first=True)
-            self.lstm = None
+            self.lstm = None  # type: ignore
 
     def update_normalization(self, vec_inputs: List[torch.Tensor]) -> None:
         for vec_input, vec_enc in zip(vec_inputs, self.vector_encoders):