Repository: instadeepai/Mava

Issue:
[INVESTIGATION]: Add layer normalisation to normalise hidden layers features
### What do you want to investigate?
The [onpolicy](https://github.com/marlbenchmark/on-policy/blob/79f6591882088a0f583f7a4bcba44041141f25f5/onpolicy/algorithms/utils/mlp.py#L19) repo implements feature normalisation by applying nn.Layernorm to the output of every linear layer. Currently we only normalise our input raw observations. We do not normlalise beyond the first layer.

### Definition of done
Add layer normalisation to every hidden layer of our networks in a similar fashion to the onpolicy repo.



Discussion:

Resolving patch:
diff --git a/examples/debugging/simple_spread/feedforward/decentralised/run_ippo_layer_norm.py b/examples/debugging/simple_spread/feedforward/decentralised/run_ippo_layer_norm.py
new file mode 100644
--- /dev/null
+++ b/examples/debugging/simple_spread/feedforward/decentralised/run_ippo_layer_norm.py
@@ -0,0 +1,121 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Example running IPPO on debug MPE environments with layer normalisation."""
+import functools
+from datetime import datetime
+from typing import Any
+
+import optax
+from absl import app, flags
+
+from mava.systems import ippo
+from mava.utils.environments import debugging_utils
+from mava.utils.loggers import logger_utils
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "env_name",
+    "simple_spread",
+    "Debugging environment name (str).",
+)
+flags.DEFINE_string(
+    "action_space",
+    "discrete",
+    "Environment action space type (str).",
+)
+
+flags.DEFINE_string(
+    "mava_id",
+    str(datetime.now()),
+    "Experiment identifier that can be used to continue experiments.",
+)
+flags.DEFINE_string("base_dir", "~/mava", "Base dir to store experiments.")
+
+
+def main(_: Any) -> None:
+    """Run main script
+
+    Args:
+        _ : _
+    """
+    # Environment.
+    environment_factory = functools.partial(
+        debugging_utils.make_environment,
+        env_name=FLAGS.env_name,
+        action_space=FLAGS.action_space,
+    )
+
+    # Networks.
+    def network_factory(*args: Any, **kwargs: Any) -> Any:
+        return ippo.make_default_networks(  # type: ignore
+            policy_layer_sizes=(64, 64),
+            critic_layer_sizes=(64, 64, 64),
+            layer_norm=True,
+            *args,
+            **kwargs,
+        )
+
+    # Used for checkpoints, tensorboard logging and env monitoring
+    experiment_path = f"{FLAGS.base_dir}/{FLAGS.mava_id}"
+
+    # Log every [log_every] seconds.
+    log_every = 10
+    logger_factory = functools.partial(
+        logger_utils.make_logger,
+        directory=FLAGS.base_dir,
+        to_terminal=True,
+        to_tensorboard=True,
+        time_stamp=FLAGS.mava_id,
+        time_delta=log_every,
+    )
+
+    # Optimisers.
+    policy_optimiser = optax.chain(
+        optax.clip_by_global_norm(40.0), optax.scale_by_adam(), optax.scale(-1e-4)
+    )
+
+    critic_optimiser = optax.chain(
+        optax.clip_by_global_norm(40.0), optax.scale_by_adam(), optax.scale(-1e-4)
+    )
+
+    # Create the system.
+    system = ippo.IPPOSystem()
+
+    # Build the system.
+    system.build(
+        environment_factory=environment_factory,
+        network_factory=network_factory,
+        logger_factory=logger_factory,
+        experiment_path=experiment_path,
+        policy_optimiser=policy_optimiser,
+        critic_optimiser=critic_optimiser,
+        run_evaluator=True,
+        sample_batch_size=5,
+        num_epochs=15,
+        num_executors=1,
+        multi_process=True,
+        clip_value=True,
+        normalize_target_values=True,
+        normalize_observations=True,
+        termination_condition={"executor_steps": 200000},
+    )
+
+    # Launch the system.
+    system.launch()
+
+
+if __name__ == "__main__":
+    app.run(main)
diff --git a/mava/systems/ippo/networks.py b/mava/systems/ippo/networks.py
--- a/mava/systems/ippo/networks.py
+++ b/mava/systems/ippo/networks.py
@@ -29,6 +29,7 @@
 from mava import specs as mava_specs
 from mava.types import NestedArray
 from mava.utils.jax_training_utils import action_mask_categorical_policies
+from mava.utils.networks_utils import MLP_NORM
 
 Array = dm_specs.Array
 BoundedArray = dm_specs.BoundedArray
@@ -204,6 +205,7 @@ def make_discrete_networks(
     orthogonal_initialisation: bool = False,
     activation_function: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
     policy_network_head_weight_gain: float = 0.01,
+    layer_norm: bool = False,
     # default behaviour is to flatten observations
 ) -> PPONetworks:
     """Create PPO network for environments with discrete action spaces.
@@ -225,6 +227,8 @@ def make_discrete_networks(
             network hidden layers.
         policy_network_head_weight_gain: value for scaling the policy
             network final layer weights by.
+        layer_norm: apply layer normalisation to the hidden MLP layers.
+
 
     Returns:
         PPONetworks class
@@ -253,11 +257,12 @@ def policy_fn(inputs: jnp.ndarray) -> networks_lib.FeedForwardNetwork:
         """
         # Add the observation network and an MLP network.
         policy_network = [
-            hk.nets.MLP(
+            MLP_NORM(
                 policy_layer_sizes,
                 activation=activation_function,
                 w_init=w_init_fn(orthogonal_initialisation, jnp.sqrt(2)),
                 activate_final=True,
+                layer_norm=layer_norm,
             ),
         ]
 
@@ -269,11 +274,12 @@ def policy_fn(inputs: jnp.ndarray) -> networks_lib.FeedForwardNetwork:
             # Add optional feedforward layers after the recurrent layers
             if len(policy_layers_after_recurrent) > 0:
                 policy_network.append(
-                    hk.nets.MLP(
+                    MLP_NORM(
                         policy_layers_after_recurrent,
                         activation=activation_function,
                         w_init=w_init_fn(orthogonal_initialisation, jnp.sqrt(2)),
                         activate_final=True,
+                        layer_norm=layer_norm,
                     ),
                 )
 
@@ -322,11 +328,12 @@ def critic_fn(inputs: jnp.ndarray) -> networks_lib.FeedForwardNetwork:
         """
         critic_network = hk.Sequential(
             [
-                hk.nets.MLP(
+                MLP_NORM(
                     critic_layer_sizes,
                     activation=activation_function,
                     w_init=w_init_fn(orthogonal_initialisation, jnp.sqrt(2)),
                     activate_final=True,
+                    layer_norm=layer_norm,
                 ),
                 ValueHead(w_init=w_init_fn(orthogonal_initialisation, 1.0)),
             ]
@@ -373,6 +380,7 @@ def make_networks(
     orthogonal_initialisation: bool = False,
     activation_function: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
     policy_network_head_weight_gain: float = 0.01,
+    layer_norm: bool = False,
 ) -> PPONetworks:
     """Function for creating PPO networks to be used.
 
@@ -396,6 +404,7 @@ def make_networks(
             network hidden layers.
         policy_network_head_weight_gain: value for scaling the policy
             network final layer weights by.
+        layer_norm: apply layer normalisation to the hidden MLP layers.
 
     Returns:
         make_discrete_networks: function to create a discrete network
@@ -415,6 +424,7 @@ def make_networks(
             recurrent_architecture_fn=recurrent_architecture_fn,
             orthogonal_initialisation=orthogonal_initialisation,
             activation_function=activation_function,
+            layer_norm=layer_norm,
             policy_network_head_weight_gain=policy_network_head_weight_gain,
         )
 
@@ -443,6 +453,7 @@ def make_default_networks(
     orthogonal_initialisation: bool = False,
     activation_function: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
     policy_network_head_weight_gain: float = 0.01,
+    layer_norm: bool = False,
 ) -> Dict[str, Any]:
     """Create default PPO networks
 
@@ -472,6 +483,7 @@ def make_default_networks(
             network hidden layers.
         policy_network_head_weight_gain: value for scaling the policy
             network final layer weights by.
+        layer_norm: apply layer normalisation to the hidden MLP layers.
 
     Returns:
         networks: networks created to given spec
@@ -498,6 +510,7 @@ def make_default_networks(
             policy_layers_after_recurrent=policy_layers_after_recurrent,
             orthogonal_initialisation=orthogonal_initialisation,
             activation_function=activation_function,
+            layer_norm=layer_norm,
             policy_network_head_weight_gain=policy_network_head_weight_gain,
         )
 
diff --git a/mava/utils/networks_utils.py b/mava/utils/networks_utils.py
new file mode 100644
--- /dev/null
+++ b/mava/utils/networks_utils.py
@@ -0,0 +1,138 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Implementation of an mlp module with layer normalisation."""
+# Adapted from https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/nets/mlp.py
+
+from typing import Any, Callable, Iterable, Optional
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+class MLP_NORM(hk.Module):
+    """A multi-layer perceptron module with an option for layer normalisation."""
+
+    def __init__(
+        self,
+        output_sizes: Iterable[int],
+        w_init: Optional[hk.initializers.Initializer] = None,
+        b_init: Optional[hk.initializers.Initializer] = None,
+        with_bias: bool = True,
+        activation: Callable[[jnp.ndarray], jnp.ndarray] = jax.nn.relu,
+        activate_final: bool = False,
+        layer_norm: bool = False,
+        name: Optional[str] = None,
+    ):
+        """Constructs an MLP with layer normalisation (MLP_NORM).
+        Args:
+            output_sizes: Sequence of layer sizes.
+            w_init: Initializer for :class:`~haiku.Linear` weights.
+            b_init: Initializer for :class:`~haiku.Linear` bias. Must be ``None`` if
+            ``with_bias=False``.
+            with_bias: Whether or not to apply a bias in each layer.
+            activation: Activation function to apply between :class:`~haiku.Linear`
+            layers. Defaults to ReLU.
+            activate_final: Whether or not to activate the final layer of the MLP.
+            layer_norm: apply layer normalisation to the hidden MLP layers.
+            name: Optional name for this module.
+        Raises:
+            ValueError: If ``with_bias`` is ``False`` and ``b_init`` is not ``None``.
+        """
+        if not with_bias and b_init is not None:
+            raise ValueError("When with_bias=False b_init must not be set.")
+
+        super().__init__(name=name)
+        self.with_bias = with_bias
+        self.w_init = w_init
+        self.b_init = b_init
+        self.activation = activation
+        self.activate_final = activate_final
+        self.layer_norm = layer_norm
+        layers = []
+        norms = []
+        output_sizes = tuple(output_sizes)
+        if self.layer_norm:
+            for index, output_size in enumerate(output_sizes):
+                layers.append(
+                    hk.Linear(
+                        output_size=output_size,
+                        w_init=w_init,
+                        b_init=b_init,
+                        with_bias=with_bias,
+                        name="linear_%d" % index,
+                    )
+                )
+                norms.append(
+                    hk.LayerNorm(
+                        axis=-1,
+                        create_scale=True,
+                        create_offset=True,
+                        param_axis=-1,
+                        use_fast_variance=False,
+                        name="norm_%d" % index,
+                    )
+                )
+            self.norms = tuple(norms)
+        else:
+            for index, output_size in enumerate(output_sizes):
+                layers.append(
+                    hk.Linear(
+                        output_size=output_size,
+                        w_init=w_init,
+                        b_init=b_init,
+                        with_bias=with_bias,
+                        name="linear_%d" % index,
+                    )
+                )
+
+        self.layers = tuple(layers)
+        self.output_size = output_sizes[-1] if output_sizes else None
+
+    def __call__(
+        self,
+        inputs: jnp.ndarray,
+        dropout_rate: Optional[float] = None,
+        rng: Any = None,
+    ) -> jnp.ndarray:
+        """Connects the module to some inputs.
+        Args:
+            inputs: A Tensor of shape ``[batch_size, input_size]``.
+            dropout_rate: Optional dropout rate.
+            rng: Optional RNG key. Require when using dropout.
+        Returns:
+            The output of the model of size ``[batch_size, output_size]``.
+        """
+        if dropout_rate is not None and rng is None:
+            raise ValueError("When using dropout an rng key must be passed.")
+        elif dropout_rate is None and rng is not None:
+            raise ValueError("RNG should only be passed when using dropout.")
+
+        rng = hk.PRNGSequence(rng) if rng is not None else None
+        num_layers = len(self.layers)
+
+        out = inputs
+        for i, layer in enumerate(self.layers):
+            out = layer(out)
+            if i < (num_layers - 1) or self.activate_final:
+                # Only perform dropout if we are activating the output.
+                if dropout_rate is not None:
+                    out = hk.dropout(next(rng), dropout_rate, out)  # type: ignore
+                out = self.activation(out)
+            if self.layer_norm:
+                out = self.norms[i](out)
+
+        return out


Test patch:
diff --git a/tests/components/building/networks_test.py b/tests/components/building/networks_test.py
--- a/tests/components/building/networks_test.py
+++ b/tests/components/building/networks_test.py
@@ -17,14 +17,19 @@
 
 from typing import Any, Callable, Dict, List, Sequence
 
+import haiku as hk
 import jax
+import jax.numpy as jnp
 import pytest
+from acme.jax import networks as networks_lib
+from acme.jax import utils
 
 from mava.components.building import DefaultNetworks
 from mava.components.building.networks import Networks, NetworksConfig
 from mava.core_jax import SystemBuilder
 from mava.specs import MAEnvironmentSpec
 from mava.systems import Builder
+from mava.utils.networks_utils import MLP_NORM
 from tests.mocks import make_fake_env_specs
 
 
@@ -317,3 +322,48 @@ def test_no_network_factory_before_build(test_builder: SystemBuilder) -> None:
         None.
     """
     assert not hasattr(test_builder.store, "network_factory")
+
+
+@hk.without_apply_rng
+@hk.transform
+def mlp(inputs: jnp.ndarray) -> networks_lib.FeedForwardNetwork:
+
+    output_sizes = [64, 64, 64]
+    mlp_network = hk.nets.MLP(output_sizes)
+
+    return mlp_network(inputs)
+
+
+@hk.without_apply_rng
+@hk.transform
+def mlp_norm(inputs: jnp.ndarray) -> networks_lib.FeedForwardNetwork:
+
+    output_sizes = [64, 64, 64]
+    mlp_network_norm = MLP_NORM(output_sizes, layer_norm=True)
+
+    return mlp_network_norm(inputs)
+
+
+@hk.without_apply_rng
+@hk.transform
+def mlp_norm_false(inputs: jnp.ndarray) -> networks_lib.FeedForwardNetwork:
+
+    output_sizes = [64, 64, 64]
+    mlp_network_norm = MLP_NORM(output_sizes, layer_norm=False)
+
+    return mlp_network_norm(inputs)
+
+
+def test_MLP_with_feature_norm() -> None:
+    """Test if layer normalisation is added correctly to the hidden layers"""
+
+    base_key = jax.random.PRNGKey(32)
+    dummy_obs = jnp.ones(16)
+    dummy_obs = utils.add_batch_dim(dummy_obs)
+    _, network_key = jax.random.split(base_key)
+    params = mlp.init(network_key, dummy_obs)  # type: ignore
+    params_norm = mlp_norm.init(network_key, dummy_obs)
+    params_norm_false = mlp_norm_false.init(network_key, dummy_obs)
+
+    assert len(params.keys()) == len(params_norm_false.keys())
+    assert len(params.keys()) * 2 == len(params_norm.keys())