# 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.

"""Trainer components for system builders."""

import abc
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Any, Callable, Dict, Optional

from mava.components.jax import Component
from mava.core_jax import SystemBuilder, SystemTrainer
from mava.utils.sort_utils import sort_str_num


class BaseTrainerInit(Component):
    @abc.abstractmethod
    def __init__(
        self,
        config: Any,
    ):
        """Initialise system init components.

        Args:
            config : a dataclass specifying the component parameters.
        """
        self.config = config

    @abc.abstractmethod
    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """Summary."""
        pass

    @abc.abstractmethod
    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """Summary."""
        pass

    @staticmethod
    def name() -> str:
        """Component name."""

        return "trainer_init"


class SingleTrainerInit(BaseTrainerInit):
    def __init__(self, config: SimpleNamespace = SimpleNamespace()):
        """Initialises a single trainer.

        Single trainer is used to train all networks.

        Args:
            config : a dataclass specifying the component parameters.
        """
        self.config = config

    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """Assigns trainers to networks for training.

        Args:
            builder : the system builder
        Raises:
            ValueError: Raises an error when trainer_networks is not
                        set to single_trainer.
        """
        unique_net_keys = builder.store.unique_net_keys

        # Setup trainer_networks

        builder.store.trainer_networks = {"trainer": unique_net_keys}

        # Get all the unique trainer network keys
        all_trainer_net_keys = []
        for trainer_nets in builder.store.trainer_networks.values():
            all_trainer_net_keys.extend(trainer_nets)
        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))

        # Check that all agent_net_keys are in trainer_networks
        assert unique_net_keys == unique_trainer_net_keys
        # Setup specs for each network
        builder.store.net_spec_keys = {}
        for i in range(len(unique_net_keys)):
            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
                i % len(builder.store.agents)
            ]

        # Setup table_network_config
        builder.store.table_network_config = {}
        for trainer_key in builder.store.trainer_networks.keys():
            most_matches = 0
            trainer_nets = builder.store.trainer_networks[trainer_key]
            for sample in builder.store.network_sampling_setup:
                matches = 0
                for entry in sample:
                    if entry in trainer_nets:
                        matches += 1
                if most_matches < matches:
                    matches = most_matches
                    builder.store.table_network_config[trainer_key] = sample

        builder.store.networks = builder.store.network_factory()

    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """_summary_"""
        # Convert network keys for the trainer.
        trainer.store.trainer_table_entry = trainer.store.table_network_config[
            trainer.store.trainer_id
        ]
        trainer.store.trainer_agents = trainer.store.agents[
            : len(trainer.store.trainer_table_entry)
        ]
        trainer.store.trainer_agent_net_keys = {
            agent: trainer.store.trainer_table_entry[a_i]
            for a_i, agent in enumerate(trainer.store.trainer_agents)
        }


class OneTrainerPerNetworkInit(BaseTrainerInit):
    def __init__(self, config: SimpleNamespace = SimpleNamespace()):
        """Initialises a multiple trainers.

        Different trainer will be dedicated to training each network.

        Args:
            config : a dataclass specifying the component parameters.
        """
        self.config = config

    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """.

        Args:
            builder : the system builder
        Raises:
            ValueError: Raises an error when trainer_networks is not
                        set to one_trainer_per_network.
        """
        unique_net_keys = builder.store.unique_net_keys

        # Setup trainer_networks
        builder.store.trainer_networks = {
            f"trainer_{i}": [unique_net_keys[i]] for i in range(len(unique_net_keys))
        }

        # Get all the unique trainer network keys
        all_trainer_net_keys = []
        for trainer_nets in builder.store.trainer_networks.values():
            all_trainer_net_keys.extend(trainer_nets)
        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))

        # Check that all agent_net_keys are in trainer_networks
        assert unique_net_keys == unique_trainer_net_keys
        # Setup specs for each network
        builder.store.net_spec_keys = {}
        for i in range(len(unique_net_keys)):
            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
                i % len(builder.store.agents)
            ]

        # Setup table_network_config
        builder.store.table_network_config = {}
        for trainer_key in builder.store.trainer_networks.keys():
            most_matches = 0
            trainer_nets = builder.store.trainer_networks[trainer_key]
            for sample in builder.store.network_sampling_setup:
                matches = 0
                for entry in sample:
                    if entry in trainer_nets:
                        matches += 1
                if most_matches < matches:
                    matches = most_matches
                    builder.store.table_network_config[trainer_key] = sample

        builder.store.networks = builder.store.network_factory()

    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """_summary_"""
        # Convert network keys for the trainer.
        trainer.store.trainer_table_entry = trainer.store.table_network_config[
            trainer.store.trainer_id
        ]
        trainer.store.trainer_agents = trainer.store.agents[
            : len(trainer.store.trainer_table_entry)
        ]
        trainer.store.trainer_agent_net_keys = {
            agent: trainer.store.trainer_table_entry[a_i]
            for a_i, agent in enumerate(trainer.store.trainer_agents)
        }


@dataclass
class CustomTrainerInitConfig:
    trainer_networks: Dict = field(default_factory=lambda: {})


class CustomTrainerInit(BaseTrainerInit):
    def __init__(self, config: CustomTrainerInitConfig = CustomTrainerInitConfig()):
        """Initialises custom trainers.

        Custom trainer network configuration can be given as a dictionary
        assigning specific trainers to specific networks.

        Args:
            config : a dataclass specifying the component parameters.
        """

        self.config = config

    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """Assigns trainers to networks for training.

        Args:
            builder : the system builder
        Raises:
            ValueError: Raises an error when trainer_networks is not
                        passed in as a dictionary.
        """
        trainer_networks = self.config.trainer_networks
        unique_net_keys = builder.store.unique_net_keys

        # Setup trainer_networks
        if not isinstance(trainer_networks, dict) or trainer_networks == {}:

            raise ValueError("trainer_networks must be a dictionary.")

        builder.store.trainer_networks = trainer_networks

        # Get all the unique trainer network keys
        all_trainer_net_keys = []
        for trainer_nets in builder.store.trainer_networks.values():
            all_trainer_net_keys.extend(trainer_nets)
        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))

        # Check that all agent_net_keys are in trainer_networks
        assert unique_net_keys == unique_trainer_net_keys
        # Setup specs for each network
        builder.store.net_spec_keys = {}
        for i in range(len(unique_net_keys)):
            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
                i % len(builder.store.agents)
            ]

        # Setup table_network_config
        builder.store.table_network_config = {}
        for trainer_key in builder.store.trainer_networks.keys():
            most_matches = 0
            trainer_nets = builder.store.trainer_networks[trainer_key]
            for sample in builder.store.network_sampling_setup:
                matches = 0
                for entry in sample:
                    if entry in trainer_nets:
                        matches += 1
                if most_matches < matches:
                    matches = most_matches
                    builder.store.table_network_config[trainer_key] = sample

        builder.store.networks = builder.store.network_factory()

    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """_summary_"""
        # Convert network keys for the trainer.
        trainer.store.trainer_table_entry = trainer.store.table_network_config[
            trainer.store.trainer_id
        ]
        trainer.store.trainer_agents = trainer.store.agents[
            : len(trainer.store.trainer_table_entry)
        ]
        trainer.store.trainer_agent_net_keys = {
            agent: trainer.store.trainer_table_entry[a_i]
            for a_i, agent in enumerate(trainer.store.trainer_agents)
        }

    @staticmethod
    def config_class() -> Optional[Callable]:
        """Config class used for component.

        Returns:
            config class/dataclass for component.
        """
        return CustomTrainerInitConfig

mava/components/jax/training/trainer.py
# 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.

"""Trainer components for system builders."""

import abc
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Any, Callable, Dict, Optional

from mava.components.jax import Component
from mava.core_jax import SystemBuilder, SystemTrainer
from mava.utils.sort_utils import sort_str_num


class BaseTrainerInit(Component):
    @abc.abstractmethod
    def __init__(
        self,
        config: Any,
    ):
        """Initialise system init components.

        Args:
            config : a dataclass specifying the component parameters.
        """
        self.config = config

    @abc.abstractmethod
    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """Summary."""
        pass

    @abc.abstractmethod
    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """Summary."""
        pass

    @staticmethod
    def name() -> str:
        """Component name."""

        return "trainer_init"


class SingleTrainerInit(BaseTrainerInit):
    def __init__(self, config: SimpleNamespace = SimpleNamespace()):
        """Initialises a single trainer.

        Single trainer is used to train all networks.

        Args:
            config : a dataclass specifying the component parameters.
        """
        self.config = config

    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """Assigns trainers to networks for training.

        Args:
            builder : the system builder
        Raises:
            ValueError: Raises an error when trainer_networks is not
                        set to single_trainer.
        """
        unique_net_keys = builder.store.unique_net_keys

        # Setup trainer_networks

        builder.store.trainer_networks = {"trainer": unique_net_keys}

        # Get all the unique trainer network keys
        all_trainer_net_keys = []
        for trainer_nets in builder.store.trainer_networks.values():
            all_trainer_net_keys.extend(trainer_nets)
        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))

        # Check that all agent_net_keys are in trainer_networks
        assert unique_net_keys == unique_trainer_net_keys
        # Setup specs for each network
        builder.store.net_spec_keys = {}
        for i in range(len(unique_net_keys)):
            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
                i % len(builder.store.agents)
            ]

        # Setup table_network_config
        builder.store.table_network_config = {}
        for trainer_key in builder.store.trainer_networks.keys():
            most_matches = 0
            trainer_nets = builder.store.trainer_networks[trainer_key]
            for sample in builder.store.network_sampling_setup:
                matches = 0
                for entry in sample:
                    if entry in trainer_nets:
                        matches += 1
                if most_matches < matches:
                    matches = most_matches
                    builder.store.table_network_config[trainer_key] = sample

        builder.store.networks = builder.store.network_factory()

    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """_summary_"""
        # Convert network keys for the trainer.
        trainer.store.trainer_table_entry = trainer.store.table_network_config[
            trainer.store.trainer_id
        ]
        trainer.store.trainer_agents = trainer.store.agents[
            : len(trainer.store.trainer_table_entry)
        ]
        trainer.store.trainer_agent_net_keys = {
            agent: trainer.store.trainer_table_entry[a_i]
            for a_i, agent in enumerate(trainer.store.trainer_agents)
        }


class OneTrainerPerNetworkInit(BaseTrainerInit):
    def __init__(self, config: SimpleNamespace = SimpleNamespace()):
        """Initialises a multiple trainers.

        Different trainer will be dedicated to training each network.

        Args:
            config : a dataclass specifying the component parameters.
        """
        self.config = config

    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """.

        Args:
            builder : the system builder
        Raises:
            ValueError: Raises an error when trainer_networks is not
                        set to one_trainer_per_network.
        """
        unique_net_keys = builder.store.unique_net_keys

        # Setup trainer_networks
        builder.store.trainer_networks = {
            f"trainer_{i}": [unique_net_keys[i]] for i in range(len(unique_net_keys))
        }

        # Get all the unique trainer network keys
        all_trainer_net_keys = []
        for trainer_nets in builder.store.trainer_networks.values():
            all_trainer_net_keys.extend(trainer_nets)
        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))

        # Check that all agent_net_keys are in trainer_networks
        assert unique_net_keys == unique_trainer_net_keys
        # Setup specs for each network
        builder.store.net_spec_keys = {}
        for i in range(len(unique_net_keys)):
            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
                i % len(builder.store.agents)
            ]

        # Setup table_network_config
        builder.store.table_network_config = {}
        for trainer_key in builder.store.trainer_networks.keys():
            most_matches = 0
            trainer_nets = builder.store.trainer_networks[trainer_key]
            for sample in builder.store.network_sampling_setup:
                matches = 0
                for entry in sample:
                    if entry in trainer_nets:
                        matches += 1
                if most_matches < matches:
                    matches = most_matches
                    builder.store.table_network_config[trainer_key] = sample

        builder.store.networks = builder.store.network_factory()

    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """_summary_"""
        # Convert network keys for the trainer.
        trainer.store.trainer_table_entry = trainer.store.table_network_config[
            trainer.store.trainer_id
        ]
        trainer.store.trainer_agents = trainer.store.agents[
            : len(trainer.store.trainer_table_entry)
        ]
        trainer.store.trainer_agent_net_keys = {
            agent: trainer.store.trainer_table_entry[a_i]
            for a_i, agent in enumerate(trainer.store.trainer_agents)
        }


@dataclass
class CustomTrainerInitConfig:
    trainer_networks: Dict = field(default_factory=lambda: {})


class CustomTrainerInit(BaseTrainerInit):
    def __init__(self, config: CustomTrainerInitConfig = CustomTrainerInitConfig()):
        """Initialises custom trainers.

        Custom trainer network configuration can be given as a dictionary
        assigning specific trainers to specific networks.

        Args:
            config : a dataclass specifying the component parameters.
        """

        self.config = config

    def on_building_init_end(self, builder: SystemBuilder) -> None:
        """Assigns trainers to networks for training.

        Args:
            builder : the system builder
        Raises:
            ValueError: Raises an error when trainer_networks is not
                        passed in as a dictionary.
        """
        trainer_networks = self.config.trainer_networks
        unique_net_keys = builder.store.unique_net_keys

        # Setup trainer_networks
        if not isinstance(trainer_networks, dict) or trainer_networks == {}:

            raise ValueError("trainer_networks must be a non-empty dictionary.")

        builder.store.trainer_networks = trainer_networks

        # Get all the unique trainer network keys
        all_trainer_net_keys = []
        for trainer_nets in builder.store.trainer_networks.values():
            all_trainer_net_keys.extend(trainer_nets)
        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))

        # Check that all agent_net_keys are in trainer_networks
        assert unique_net_keys == unique_trainer_net_keys
        # Setup specs for each network
        builder.store.net_spec_keys = {}
        for i in range(len(unique_net_keys)):
            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
                i % len(builder.store.agents)
            ]

        # Setup table_network_config
        builder.store.table_network_config = {}
        for trainer_key in builder.store.trainer_networks.keys():
            most_matches = 0
            trainer_nets = builder.store.trainer_networks[trainer_key]
            for sample in builder.store.network_sampling_setup:
                matches = 0
                for entry in sample:
                    if entry in trainer_nets:
                        matches += 1
                if most_matches < matches:
                    matches = most_matches
                    builder.store.table_network_config[trainer_key] = sample

        builder.store.networks = builder.store.network_factory()

    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
        """_summary_"""
        # Convert network keys for the trainer.
        trainer.store.trainer_table_entry = trainer.store.table_network_config[
            trainer.store.trainer_id
        ]
        trainer.store.trainer_agents = trainer.store.agents[
            : len(trainer.store.trainer_table_entry)
        ]
        trainer.store.trainer_agent_net_keys = {
            agent: trainer.store.trainer_table_entry[a_i]
            for a_i, agent in enumerate(trainer.store.trainer_agents)
        }

    @staticmethod
    def config_class() -> Optional[Callable]:
        """Config class used for component.

        Returns:
            config class/dataclass for component.
        """
        return CustomTrainerInitConfig

mava/components/jax/training/trainer.py
diff --git a/mava/components/jax/training/trainer.py b/mava/components/jax/training/trainer.py
index ecf0923a8..95c89ce69 100644
--- a/mava/components/jax/training/trainer.py
+++ b/mava/components/jax/training/trainer.py
@@ -233,7 +233,7 @@ def on_building_init_end(self, builder: SystemBuilder) -> None:
         # Setup trainer_networks
         if not isinstance(trainer_networks, dict) or trainer_networks == {}:
 
-            raise ValueError("trainer_networks must be a dictionary.")
+            raise ValueError("trainer_networks must be a non-empty dictionary.")
 
         builder.store.trainer_networks = trainer_networks
 
diff --git a/tests/jax/components/training/trainer_test.py b/tests/jax/components/training/trainer_test.py
new file mode 100644
index 000000000..233d60599
--- /dev/null
+++ b/tests/jax/components/training/trainer_test.py
@@ -0,0 +1,719 @@
+# python3
+# Copyright 2022 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.
+
+from types import SimpleNamespace
+from typing import List
+
+import pytest
+
+from mava.callbacks.base import Callback
+from mava.components.jax.training.trainer import (
+    CustomTrainerInit,
+    CustomTrainerInitConfig,
+    OneTrainerPerNetworkInit,
+    SingleTrainerInit,
+)
+from mava.systems.jax.builder import Builder
+from mava.systems.jax.trainer import Trainer
+
+#######################
+# TrainerInit Fixtures
+#######################
+
+
+@pytest.fixture
+def single_trainer_init() -> SingleTrainerInit:
+    """Creates single trainer init component."""
+    return SingleTrainerInit()
+
+
+@pytest.fixture
+def one_trainer_per_network_init() -> OneTrainerPerNetworkInit:
+    """Create one trainer per network init component"""
+    return OneTrainerPerNetworkInit()
+
+
+#################
+# Mock builders
+#################
+
+
+class MockBuilder(Builder):
+    def __init__(
+        self,
+        components: List[Callback],
+        global_config: SimpleNamespace = SimpleNamespace(),
+    ) -> None:
+        """Creates a mock builder for testing
+
+        Args:
+            components : list of mava components
+            global_config : global system config
+        """
+        super().__init__(components, global_config)
+        self.store.agents = ["agent_0", "agent_1", "agent_2"]
+        self.store.network_factory = lambda: "network_initialized"
+
+
+@pytest.fixture
+def mock_builder_shared_weights_fixed_sampling() -> Builder:
+    """Creates mock builder.
+
+    Has shared weights and fixed agent trainer sampling setup.
+    """
+
+    mock_builder = MockBuilder(global_config=SimpleNamespace(), components=[])
+
+    mock_builder.store.unique_net_keys = ["network_agent"]
+    mock_builder.store.network_sampling_setup = [
+        ["network_agent", "network_agent", "network_agent"]
+    ]
+
+    return mock_builder
+
+
+@pytest.fixture
+def mock_builder_no_shared_weights_fixed_sampling() -> Builder:
+    """Create mock builder.
+
+    Doesn't shared weights and has fixed agent trainer sampling setup.
+    """
+
+    mock_builder = MockBuilder(global_config=SimpleNamespace(), components=[])
+
+    mock_builder.store.unique_net_keys = [
+        "network_agent_0",
+        "network_agent_1",
+        "network_agent_2",
+    ]
+    mock_builder.store.network_sampling_setup = [
+        ["network_agent_0", "network_agent_1", "network_agent_2"]
+    ]
+
+    return mock_builder
+
+
+@pytest.fixture
+def mock_builder_no_shared_weights_random_sampling() -> Builder:
+    """Create mock builder.
+
+    Doesn't have shared weights and has random agent trainer sampling setup
+    """
+
+    mock_builder = MockBuilder(global_config=SimpleNamespace(), components=[])
+
+    mock_builder.store.unique_net_keys = ["network_0", "network_1", "network_2"]
+    mock_builder.store.network_sampling_setup = [
+        ["network_0"],
+        ["network_1"],
+        ["network_2"],
+    ]
+
+    return mock_builder
+
+
+##################
+# Mock trainers
+##################
+
+
+class MockTrainer(Trainer):
+    def __init__(
+        self,
+        store: SimpleNamespace,
+        components: List[Callback] = [],
+    ) -> None:
+        """Creates mock trainer for testing
+
+        Args:
+            store : trainer store
+            components : list of Mava components
+        """
+
+        super().__init__(store, components)
+        self.store.agents = ["agent_0", "agent_1", "agent_2"]
+
+
+@pytest.fixture
+def mock_single_trainer_shared_weights_fixed_sampling() -> Trainer:
+    """Creates mock trainer
+
+    Has shared weights and fixed agent trainer sampling setup.
+    """
+
+    mock_trainer = MockTrainer(store=SimpleNamespace(), components=[])
+
+    mock_trainer.store.trainer_id = "trainer"
+    mock_trainer.store.table_network_config = {
+        "trainer": ["network_agent", "network_agent", "network_agent"]
+    }
+
+    return mock_trainer
+
+
+@pytest.fixture
+def mock_single_trainer_no_shared_weights_fixed_sampling() -> Trainer:
+    """Creates mock trainer
+
+    Doesn't have shared weights and has fixed agent
+    trainer sampling setup.
+    """
+
+    mock_trainer = MockTrainer(store=SimpleNamespace(), components=[])
+
+    mock_trainer.store.trainer_id = "trainer"
+    mock_trainer.store.table_network_config = {
+        "trainer": ["network_agent_0", "network_agent_1", "network_agent_2"]
+    }
+
+    return mock_trainer
+
+
+@pytest.fixture
+def mock_single_trainer_no_shared_weights_random_sampling() -> Trainer:
+    """Creates mock trainer
+
+    Doesn't have shared weights and has random agent
+    trainer sampling setup.
+    """
+
+    mock_trainer = MockTrainer(store=SimpleNamespace(), components=[])
+
+    mock_trainer.store.trainer_id = "trainer"
+    mock_trainer.store.table_network_config = {"trainer": ["network_2"]}
+
+    return mock_trainer
+
+
+@pytest.fixture
+def mock_one_trainer_per_network_shared_weights_fixed_sampling() -> Trainer:
+    """Creates mock trainer
+
+    Doesn't have shared weights and has fixed agent
+    trainer sampling setup.
+    """
+
+    mock_trainer = MockTrainer(store=SimpleNamespace(), components=[])
+
+    mock_trainer.store.trainer_id = "trainer_0"
+    mock_trainer.store.table_network_config = {
+        "trainer_0": ["network_agent", "network_agent", "network_agent"]
+    }
+
+    return mock_trainer
+
+
+@pytest.fixture
+def mock_one_trainer_per_network_no_shared_weights_fixed_sampling() -> Trainer:
+    """Creates mock trainer
+
+    Doesn't have shared weights and has fixed agent
+    trainer sampling setup.
+    """
+
+    mock_trainer = MockTrainer(store=SimpleNamespace(), components=[])
+
+    # trainer id must be specified per test
+    mock_trainer.store.table_network_config = {
+        "trainer_0": ["network_agent_0", "network_agent_1", "network_agent_2"],
+        "trainer_1": ["network_agent_0", "network_agent_1", "network_agent_2"],
+        "trainer_2": ["network_agent_0", "network_agent_1", "network_agent_2"],
+    }
+
+    return mock_trainer
+
+
+@pytest.fixture
+def mock_one_trainer_per_network_random_sampling() -> Trainer:
+    """Creates mock trainer
+
+    Doesn't have shared weights and has random agent
+    trainer sampling setup.
+    """
+
+    mock_trainer = MockTrainer(store=SimpleNamespace(), components=[])
+
+    # trainer id must be given at each test
+    mock_trainer.store.table_network_config = {
+        "trainer_0": ["network_0"],
+        "trainer_1": ["network_1"],
+        "trainer_2": ["network_2"],
+    }
+
+    return mock_trainer
+
+
+############################
+# TESTS
+############################
+
+# ON_BUILDING_INIT_END TESTS
+
+
+def test_single_trainer_shared_weights_fixed_sampling(
+    mock_builder_shared_weights_fixed_sampling: Builder,
+    single_trainer_init: SingleTrainerInit,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    Single trainer for all networks, shared network weights and
+    fixed agent network sampling.
+    """
+
+    builder = mock_builder_shared_weights_fixed_sampling
+    single_trainer_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {"network_agent": "agent_0"}
+    assert builder.store.table_network_config == {
+        "trainer": ["network_agent", "network_agent", "network_agent"]
+    }
+    assert builder.store.trainer_networks == {"trainer": ["network_agent"]}
+    assert builder.store.networks == "network_initialized"
+
+
+def test_single_trainer_no_shared_weights_fixed_sampling(
+    mock_builder_no_shared_weights_fixed_sampling: Builder,
+    single_trainer_init: SingleTrainerInit,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    Single trainer for all networks, no shared network weights and
+    fixed agent network sampling.
+    """
+
+    builder = mock_builder_no_shared_weights_fixed_sampling
+    single_trainer_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {
+        "network_agent_0": "agent_0",
+        "network_agent_1": "agent_1",
+        "network_agent_2": "agent_2",
+    }
+    assert builder.store.table_network_config == {
+        "trainer": ["network_agent_0", "network_agent_1", "network_agent_2"]
+    }
+    assert builder.store.trainer_networks == {
+        "trainer": ["network_agent_0", "network_agent_1", "network_agent_2"]
+    }
+    assert builder.store.networks == "network_initialized"
+
+
+def test_single_trainer_no_shared_weights_random_sampling(
+    mock_builder_no_shared_weights_random_sampling: Builder,
+    single_trainer_init: SingleTrainerInit,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    Single trainer for all networks, no shared network weights and
+    random agent network sampling.
+    """
+
+    builder = mock_builder_no_shared_weights_random_sampling
+    single_trainer_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {
+        "network_0": "agent_0",
+        "network_1": "agent_1",
+        "network_2": "agent_2",
+    }
+    assert builder.store.table_network_config == {"trainer": ["network_2"]}
+    assert builder.store.trainer_networks == {
+        "trainer": ["network_0", "network_1", "network_2"]
+    }
+    assert builder.store.networks == "network_initialized"
+
+
+def test_one_trainer_per_network_shared_weights_fixed_sampling(
+    mock_builder_shared_weights_fixed_sampling: Builder,
+    one_trainer_per_network_init: OneTrainerPerNetworkInit,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    One trainer per network, shared network weights and
+    fixed agent network sampling.
+    """
+
+    builder = mock_builder_shared_weights_fixed_sampling
+    one_trainer_per_network_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {"network_agent": "agent_0"}
+    assert builder.store.table_network_config == {
+        "trainer_0": ["network_agent", "network_agent", "network_agent"]
+    }
+    assert builder.store.trainer_networks == {"trainer_0": ["network_agent"]}
+    assert builder.store.networks == "network_initialized"
+
+
+def test_one_trainer_per_network_no_shared_weights_fixed_sampling(
+    mock_builder_no_shared_weights_fixed_sampling: Builder,
+    one_trainer_per_network_init: OneTrainerPerNetworkInit,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    One trainer per network, no shared network weights and
+    fixed agent network sampling.
+    """
+
+    builder = mock_builder_no_shared_weights_fixed_sampling
+    one_trainer_per_network_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {
+        "network_agent_0": "agent_0",
+        "network_agent_1": "agent_1",
+        "network_agent_2": "agent_2",
+    }
+    assert builder.store.table_network_config == {
+        "trainer_0": ["network_agent_0", "network_agent_1", "network_agent_2"],
+        "trainer_1": ["network_agent_0", "network_agent_1", "network_agent_2"],
+        "trainer_2": ["network_agent_0", "network_agent_1", "network_agent_2"],
+    }
+    assert builder.store.trainer_networks == {
+        "trainer_0": ["network_agent_0"],
+        "trainer_1": ["network_agent_1"],
+        "trainer_2": ["network_agent_2"],
+    }
+    assert builder.store.networks == "network_initialized"
+
+
+def test_one_trainer_per_network_random_sampling(
+    mock_builder_no_shared_weights_random_sampling: Builder,
+    one_trainer_per_network_init: OneTrainerPerNetworkInit,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    One trainer per network, no shared network weights and
+    random agent network sampling.
+    """
+
+    builder = mock_builder_no_shared_weights_random_sampling
+    one_trainer_per_network_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {
+        "network_0": "agent_0",
+        "network_1": "agent_1",
+        "network_2": "agent_2",
+    }
+    assert builder.store.table_network_config == {
+        "trainer_0": ["network_0"],
+        "trainer_1": ["network_1"],
+        "trainer_2": ["network_2"],
+    }
+    assert builder.store.trainer_networks == {
+        "trainer_0": ["network_0"],
+        "trainer_1": ["network_1"],
+        "trainer_2": ["network_2"],
+    }
+    assert builder.store.networks == "network_initialized"
+
+
+def test_custom_trainer_on_building_init_value_error(
+    mock_builder_no_shared_weights_fixed_sampling: Builder,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    Custom trainer init. Asserts than exception is thrown when an empty dictionary
+    is passed in.
+    """
+
+    trainer_init = CustomTrainerInit(config=CustomTrainerInitConfig())
+
+    builder = mock_builder_no_shared_weights_fixed_sampling
+
+    with pytest.raises(ValueError):
+        trainer_init.on_building_init_end(builder)
+
+    trainer_init = CustomTrainerInit(
+        config=CustomTrainerInitConfig(trainer_networks=None)  # type:ignore
+    )
+
+    builder = mock_builder_no_shared_weights_fixed_sampling
+
+    with pytest.raises(ValueError):
+        trainer_init.on_building_init_end(builder)
+
+
+def test_custom_trainer_init_no_shared_weights_random_sampling(
+    mock_builder_no_shared_weights_random_sampling: Builder,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    Custom trainer init. Asserts that custom trainer init works correctly
+    when trainer networks corresponding to random network sampling and no
+    shared weights are passed in.
+    """
+
+    trainer_init = CustomTrainerInit(
+        config=CustomTrainerInitConfig(
+            trainer_networks={
+                "trainer_0": ["network_0"],
+                "trainer_1": ["network_1"],
+                "trainer_2": ["network_2"],
+            }
+        )
+    )
+
+    builder = mock_builder_no_shared_weights_random_sampling
+
+    trainer_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {
+        "network_0": "agent_0",
+        "network_1": "agent_1",
+        "network_2": "agent_2",
+    }
+
+    assert builder.store.table_network_config == {
+        "trainer_0": ["network_0"],
+        "trainer_1": ["network_1"],
+        "trainer_2": ["network_2"],
+    }
+
+    assert builder.store.networks == "network_initialized"
+
+
+def test_custom_trainer_init_shared_weights_fixed_sampling(
+    mock_builder_shared_weights_fixed_sampling: Builder,
+) -> None:
+    """Tests on_building_init_end hook.
+
+    Custom trainer init. Asserts that custom trainer init works correctly
+    when trainer networks corresponding to fixed network sampling and shared
+    network weights are passed in.
+    """
+
+    trainer_init = CustomTrainerInit(
+        config=CustomTrainerInitConfig(trainer_networks={"trainer": ["network_agent"]})
+    )
+
+    builder = mock_builder_shared_weights_fixed_sampling
+
+    trainer_init.on_building_init_end(builder)
+
+    assert builder.store.net_spec_keys == {"network_agent": "agent_0"}
+    assert builder.store.table_network_config == {
+        "trainer": ["network_agent", "network_agent", "network_agent"]
+    }
+    assert builder.store.networks == "network_initialized"
+
+
+#################################
+# ON_TRAINING_UTILITY_FNS TESTS
+#################################
+
+
+def test_training_utility_single_trainer_shared_weights(
+    mock_single_trainer_shared_weights_fixed_sampling: Trainer,
+    single_trainer_init: SingleTrainerInit,
+) -> None:
+    """Tests on_training_utility_fn hook in TrainerInit
+
+    Tests with single trainer for all agents where network weights are
+    shared and the trainer sampling setup is fixed.
+    """
+
+    trainer = mock_single_trainer_shared_weights_fixed_sampling
+
+    single_trainer_init.on_training_utility_fns(trainer)
+
+    assert trainer.store.trainer_table_entry == [
+        "network_agent",
+        "network_agent",
+        "network_agent",
+    ]
+    assert trainer.store.trainer_agents == ["agent_0", "agent_1", "agent_2"]
+    assert trainer.store.trainer_agent_net_keys == {
+        "agent_0": "network_agent",
+        "agent_1": "network_agent",
+        "agent_2": "network_agent",
+    }
+
+
+def test_training_utility_single_trainer_no_shared_weights(
+    mock_single_trainer_no_shared_weights_fixed_sampling: Trainer,
+    single_trainer_init: SingleTrainerInit,
+) -> None:
+    """Tests on_training_utility_fn hook in TrainerInit
+
+    Tests with single trainer for all agents where network weights aren't
+    shared and the trainer sampling setup is fixed.
+    """
+
+    trainer = mock_single_trainer_no_shared_weights_fixed_sampling
+
+    single_trainer_init.on_training_utility_fns(trainer)
+
+    assert trainer.store.trainer_table_entry == [
+        "network_agent_0",
+        "network_agent_1",
+        "network_agent_2",
+    ]
+    assert trainer.store.trainer_agents == ["agent_0", "agent_1", "agent_2"]
+    assert trainer.store.trainer_agent_net_keys == {
+        "agent_0": "network_agent_0",
+        "agent_1": "network_agent_1",
+        "agent_2": "network_agent_2",
+    }
+
+
+def test_training_utility_single_trainer_no_shared_weights_random_sampling(
+    mock_single_trainer_no_shared_weights_random_sampling: Trainer,
+    single_trainer_init: SingleTrainerInit,
+) -> None:
+    """Tests on_training_utility_fn hook in TrainerInit
+
+    Tests with single trainer for all agents where network weights aren't
+    shared and the trainer sampling setup is random.
+    """
+
+    trainer = mock_single_trainer_no_shared_weights_random_sampling
+
+    single_trainer_init.on_training_utility_fns(trainer)
+
+    assert trainer.store.trainer_table_entry == ["network_2"]
+    assert trainer.store.trainer_agents == ["agent_0"]
+    assert trainer.store.trainer_agent_net_keys == {"agent_0": "network_2"}
+
+
+def test_training_utility_one_trainer_per_network_shared_weights(
+    mock_one_trainer_per_network_shared_weights_fixed_sampling: Trainer,
+    one_trainer_per_network_init: OneTrainerPerNetworkInit,
+) -> None:
+    """Tests on_training_utility_fns hook.
+
+        Tests with one trainer per network where network weights are
+    shared and the trainer sampling setup is fixed.
+    """
+
+    trainer = mock_one_trainer_per_network_shared_weights_fixed_sampling
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer)
+
+    assert trainer.store.trainer_table_entry == [
+        "network_agent",
+        "network_agent",
+        "network_agent",
+    ]
+    assert trainer.store.trainer_agents == ["agent_0", "agent_1", "agent_2"]
+    assert trainer.store.trainer_agent_net_keys == {
+        "agent_0": "network_agent",
+        "agent_1": "network_agent",
+        "agent_2": "network_agent",
+    }
+
+
+def test_training_utility_one_trainer_per_network_no_shared_weights(
+    mock_one_trainer_per_network_no_shared_weights_fixed_sampling: Trainer,
+    one_trainer_per_network_init: OneTrainerPerNetworkInit,
+) -> None:
+    """Tests on_training_utility_fns hook.
+
+        Tests with one trainer per network where network weights aren't
+    shared and the trainer sampling setup is fixed.
+    """
+
+    trainer_0 = mock_one_trainer_per_network_no_shared_weights_fixed_sampling
+
+    trainer_0.store.trainer_id = "trainer_0"
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer_0)
+
+    assert trainer_0.store.trainer_table_entry == [
+        "network_agent_0",
+        "network_agent_1",
+        "network_agent_2",
+    ]
+    assert trainer_0.store.trainer_agents == ["agent_0", "agent_1", "agent_2"]
+    assert trainer_0.store.trainer_agent_net_keys == {
+        "agent_0": "network_agent_0",
+        "agent_1": "network_agent_1",
+        "agent_2": "network_agent_2",
+    }
+
+    trainer_1 = mock_one_trainer_per_network_no_shared_weights_fixed_sampling
+
+    trainer_1.store.trainer_id = "trainer_1"
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer_1)
+
+    assert trainer_1.store.trainer_table_entry == [
+        "network_agent_0",
+        "network_agent_1",
+        "network_agent_2",
+    ]
+    assert trainer_1.store.trainer_agents == ["agent_0", "agent_1", "agent_2"]
+    assert trainer_1.store.trainer_agent_net_keys == {
+        "agent_0": "network_agent_0",
+        "agent_1": "network_agent_1",
+        "agent_2": "network_agent_2",
+    }
+
+    trainer_2 = mock_one_trainer_per_network_no_shared_weights_fixed_sampling
+
+    trainer_2.store.trainer_id = "trainer_2"
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer_2)
+
+    assert trainer_2.store.trainer_table_entry == [
+        "network_agent_0",
+        "network_agent_1",
+        "network_agent_2",
+    ]
+    assert trainer_2.store.trainer_agents == ["agent_0", "agent_1", "agent_2"]
+    assert trainer_2.store.trainer_agent_net_keys == {
+        "agent_0": "network_agent_0",
+        "agent_1": "network_agent_1",
+        "agent_2": "network_agent_2",
+    }
+
+
+def test_training_utility_one_trainer_per_network_random_sampling(
+    mock_one_trainer_per_network_random_sampling: Trainer,
+    one_trainer_per_network_init: OneTrainerPerNetworkInit,
+) -> None:
+    """Tests on_training_utility_fns hook.
+
+        Tests with one trainer per network where network weights aren't
+    shared and the trainer sampling setup is random.
+    """
+
+    trainer_0 = mock_one_trainer_per_network_random_sampling
+
+    trainer_0.store.trainer_id = "trainer_0"
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer_0)
+
+    assert trainer_0.store.trainer_table_entry == ["network_0"]
+    assert trainer_0.store.trainer_agents == ["agent_0"]
+    assert trainer_0.store.trainer_agent_net_keys == {"agent_0": "network_0"}
+
+    trainer_1 = mock_one_trainer_per_network_random_sampling
+
+    trainer_1.store.trainer_id = "trainer_1"
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer_1)
+
+    assert trainer_1.store.trainer_table_entry == ["network_1"]
+    assert trainer_1.store.trainer_agents == ["agent_0"]
+    assert trainer_1.store.trainer_agent_net_keys == {"agent_0": "network_1"}
+
+    trainer_2 = mock_one_trainer_per_network_random_sampling
+
+    trainer_2.store.trainer_id = "trainer_2"
+
+    one_trainer_per_network_init.on_training_utility_fns(trainer_2)
+
+    assert trainer_2.store.trainer_table_entry == ["network_2"]
+    assert trainer_2.store.trainer_agents == ["agent_0"]
+    assert trainer_2.store.trainer_agent_net_keys == {"agent_0": "network_2"}

[TEST] Jax Trainer Init
### What do you want to test?
Jax trainer init components

### Outline of test structure
* Unit tests
* Test components and hooks

### Definition of done
Passing checks, cover all hooks, edge cases considered

### Mandatory checklist before making a PR
* [ ] The success criteria laid down in “Definition of done” are met.
* [ ] Test code is documented - docstrings for methods and classes, static types for arguments.
* [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.
