We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mava/components/jax/training/trainer.py`
Content:
```
1 # python3
2 # Copyright 2021 InstaDeep Ltd. All rights reserved.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 
16 """Trainer components for system builders."""
17 
18 import abc
19 from dataclasses import dataclass, field
20 from types import SimpleNamespace
21 from typing import Any, Callable, Dict, Optional
22 
23 from mava.components.jax import Component
24 from mava.core_jax import SystemBuilder, SystemTrainer
25 from mava.utils.sort_utils import sort_str_num
26 
27 
28 class BaseTrainerInit(Component):
29     @abc.abstractmethod
30     def __init__(
31         self,
32         config: Any,
33     ):
34         """Initialise system init components.
35 
36         Args:
37             config : a dataclass specifying the component parameters.
38         """
39         self.config = config
40 
41     @abc.abstractmethod
42     def on_building_init_end(self, builder: SystemBuilder) -> None:
43         """Summary."""
44         pass
45 
46     @abc.abstractmethod
47     def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
48         """Summary."""
49         pass
50 
51     @staticmethod
52     def name() -> str:
53         """Component name."""
54 
55         return "trainer_init"
56 
57 
58 class SingleTrainerInit(BaseTrainerInit):
59     def __init__(self, config: SimpleNamespace = SimpleNamespace()):
60         """Initialises a single trainer.
61 
62         Single trainer is used to train all networks.
63 
64         Args:
65             config : a dataclass specifying the component parameters.
66         """
67         self.config = config
68 
69     def on_building_init_end(self, builder: SystemBuilder) -> None:
70         """Assigns trainers to networks for training.
71 
72         Args:
73             builder : the system builder
74         Raises:
75             ValueError: Raises an error when trainer_networks is not
76                         set to single_trainer.
77         """
78         unique_net_keys = builder.store.unique_net_keys
79 
80         # Setup trainer_networks
81 
82         builder.store.trainer_networks = {"trainer": unique_net_keys}
83 
84         # Get all the unique trainer network keys
85         all_trainer_net_keys = []
86         for trainer_nets in builder.store.trainer_networks.values():
87             all_trainer_net_keys.extend(trainer_nets)
88         unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))
89 
90         # Check that all agent_net_keys are in trainer_networks
91         assert unique_net_keys == unique_trainer_net_keys
92         # Setup specs for each network
93         builder.store.net_spec_keys = {}
94         for i in range(len(unique_net_keys)):
95             builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
96                 i % len(builder.store.agents)
97             ]
98 
99         # Setup table_network_config
100         builder.store.table_network_config = {}
101         for trainer_key in builder.store.trainer_networks.keys():
102             most_matches = 0
103             trainer_nets = builder.store.trainer_networks[trainer_key]
104             for sample in builder.store.network_sampling_setup:
105                 matches = 0
106                 for entry in sample:
107                     if entry in trainer_nets:
108                         matches += 1
109                 if most_matches < matches:
110                     matches = most_matches
111                     builder.store.table_network_config[trainer_key] = sample
112 
113         builder.store.networks = builder.store.network_factory()
114 
115     def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
116         """_summary_"""
117         # Convert network keys for the trainer.
118         trainer.store.trainer_table_entry = trainer.store.table_network_config[
119             trainer.store.trainer_id
120         ]
121         trainer.store.trainer_agents = trainer.store.agents[
122             : len(trainer.store.trainer_table_entry)
123         ]
124         trainer.store.trainer_agent_net_keys = {
125             agent: trainer.store.trainer_table_entry[a_i]
126             for a_i, agent in enumerate(trainer.store.trainer_agents)
127         }
128 
129 
130 class OneTrainerPerNetworkInit(BaseTrainerInit):
131     def __init__(self, config: SimpleNamespace = SimpleNamespace()):
132         """Initialises a multiple trainers.
133 
134         Different trainer will be dedicated to training each network.
135 
136         Args:
137             config : a dataclass specifying the component parameters.
138         """
139         self.config = config
140 
141     def on_building_init_end(self, builder: SystemBuilder) -> None:
142         """.
143 
144         Args:
145             builder : the system builder
146         Raises:
147             ValueError: Raises an error when trainer_networks is not
148                         set to one_trainer_per_network.
149         """
150         unique_net_keys = builder.store.unique_net_keys
151 
152         # Setup trainer_networks
153         builder.store.trainer_networks = {
154             f"trainer_{i}": [unique_net_keys[i]] for i in range(len(unique_net_keys))
155         }
156 
157         # Get all the unique trainer network keys
158         all_trainer_net_keys = []
159         for trainer_nets in builder.store.trainer_networks.values():
160             all_trainer_net_keys.extend(trainer_nets)
161         unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))
162 
163         # Check that all agent_net_keys are in trainer_networks
164         assert unique_net_keys == unique_trainer_net_keys
165         # Setup specs for each network
166         builder.store.net_spec_keys = {}
167         for i in range(len(unique_net_keys)):
168             builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
169                 i % len(builder.store.agents)
170             ]
171 
172         # Setup table_network_config
173         builder.store.table_network_config = {}
174         for trainer_key in builder.store.trainer_networks.keys():
175             most_matches = 0
176             trainer_nets = builder.store.trainer_networks[trainer_key]
177             for sample in builder.store.network_sampling_setup:
178                 matches = 0
179                 for entry in sample:
180                     if entry in trainer_nets:
181                         matches += 1
182                 if most_matches < matches:
183                     matches = most_matches
184                     builder.store.table_network_config[trainer_key] = sample
185 
186         builder.store.networks = builder.store.network_factory()
187 
188     def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
189         """_summary_"""
190         # Convert network keys for the trainer.
191         trainer.store.trainer_table_entry = trainer.store.table_network_config[
192             trainer.store.trainer_id
193         ]
194         trainer.store.trainer_agents = trainer.store.agents[
195             : len(trainer.store.trainer_table_entry)
196         ]
197         trainer.store.trainer_agent_net_keys = {
198             agent: trainer.store.trainer_table_entry[a_i]
199             for a_i, agent in enumerate(trainer.store.trainer_agents)
200         }
201 
202 
203 @dataclass
204 class CustomTrainerInitConfig:
205     trainer_networks: Dict = field(default_factory=lambda: {})
206 
207 
208 class CustomTrainerInit(BaseTrainerInit):
209     def __init__(self, config: CustomTrainerInitConfig = CustomTrainerInitConfig()):
210         """Initialises custom trainers.
211 
212         Custom trainer network configuration can be given as a dictionary
213         assigning specific trainers to specific networks.
214 
215         Args:
216             config : a dataclass specifying the component parameters.
217         """
218 
219         self.config = config
220 
221     def on_building_init_end(self, builder: SystemBuilder) -> None:
222         """Assigns trainers to networks for training.
223 
224         Args:
225             builder : the system builder
226         Raises:
227             ValueError: Raises an error when trainer_networks is not
228                         passed in as a dictionary.
229         """
230         trainer_networks = self.config.trainer_networks
231         unique_net_keys = builder.store.unique_net_keys
232 
233         # Setup trainer_networks
234         if not isinstance(trainer_networks, dict) or trainer_networks == {}:
235 
236             raise ValueError("trainer_networks must be a dictionary.")
237 
238         builder.store.trainer_networks = trainer_networks
239 
240         # Get all the unique trainer network keys
241         all_trainer_net_keys = []
242         for trainer_nets in builder.store.trainer_networks.values():
243             all_trainer_net_keys.extend(trainer_nets)
244         unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))
245 
246         # Check that all agent_net_keys are in trainer_networks
247         assert unique_net_keys == unique_trainer_net_keys
248         # Setup specs for each network
249         builder.store.net_spec_keys = {}
250         for i in range(len(unique_net_keys)):
251             builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[
252                 i % len(builder.store.agents)
253             ]
254 
255         # Setup table_network_config
256         builder.store.table_network_config = {}
257         for trainer_key in builder.store.trainer_networks.keys():
258             most_matches = 0
259             trainer_nets = builder.store.trainer_networks[trainer_key]
260             for sample in builder.store.network_sampling_setup:
261                 matches = 0
262                 for entry in sample:
263                     if entry in trainer_nets:
264                         matches += 1
265                 if most_matches < matches:
266                     matches = most_matches
267                     builder.store.table_network_config[trainer_key] = sample
268 
269         builder.store.networks = builder.store.network_factory()
270 
271     def on_training_utility_fns(self, trainer: SystemTrainer) -> None:
272         """_summary_"""
273         # Convert network keys for the trainer.
274         trainer.store.trainer_table_entry = trainer.store.table_network_config[
275             trainer.store.trainer_id
276         ]
277         trainer.store.trainer_agents = trainer.store.agents[
278             : len(trainer.store.trainer_table_entry)
279         ]
280         trainer.store.trainer_agent_net_keys = {
281             agent: trainer.store.trainer_table_entry[a_i]
282             for a_i, agent in enumerate(trainer.store.trainer_agents)
283         }
284 
285     @staticmethod
286     def config_class() -> Optional[Callable]:
287         """Config class used for component.
288 
289         Returns:
290             config class/dataclass for component.
291         """
292         return CustomTrainerInitConfig
293
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mava/components/jax/training/trainer.py b/mava/components/jax/training/trainer.py
--- a/mava/components/jax/training/trainer.py
+++ b/mava/components/jax/training/trainer.py
@@ -233,7 +233,7 @@
         # 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

{"golden_diff": "diff --git a/mava/components/jax/training/trainer.py b/mava/components/jax/training/trainer.py\n--- a/mava/components/jax/training/trainer.py\n+++ b/mava/components/jax/training/trainer.py\n@@ -233,7 +233,7 @@\n         # Setup trainer_networks\n         if not isinstance(trainer_networks, dict) or trainer_networks == {}:\n \n-            raise ValueError(\"trainer_networks must be a dictionary.\")\n+            raise ValueError(\"trainer_networks must be a non-empty dictionary.\")\n \n         builder.store.trainer_networks = trainer_networks\n", "issue": "[TEST] Jax Trainer Init\n### What do you want to test?\r\nJax trainer init components\r\n\r\n### Outline of test structure\r\n* Unit tests\r\n* Test components and hooks\r\n\r\n### Definition of done\r\nPassing checks, cover all hooks, edge cases considered\r\n\r\n### Mandatory checklist before making a PR\r\n* [ ] The success criteria laid down in \u201cDefinition of done\u201d are met.\r\n* [ ] Test code is documented - docstrings for methods and classes, static types for arguments.\r\n* [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.\n", "before_files": [{"content": "# python3\n# Copyright 2021 InstaDeep Ltd. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Trainer components for system builders.\"\"\"\n\nimport abc\nfrom dataclasses import dataclass, field\nfrom types import SimpleNamespace\nfrom typing import Any, Callable, Dict, Optional\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemBuilder, SystemTrainer\nfrom mava.utils.sort_utils import sort_str_num\n\n\nclass BaseTrainerInit(Component):\n    @abc.abstractmethod\n    def __init__(\n        self,\n        config: Any,\n    ):\n        \"\"\"Initialise system init components.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    @abc.abstractmethod\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\"Summary.\"\"\"\n        pass\n\n    @abc.abstractmethod\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"Summary.\"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"Component name.\"\"\"\n\n        return \"trainer_init\"\n\n\nclass SingleTrainerInit(BaseTrainerInit):\n    def __init__(self, config: SimpleNamespace = SimpleNamespace()):\n        \"\"\"Initialises a single trainer.\n\n        Single trainer is used to train all networks.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\"Assigns trainers to networks for training.\n\n        Args:\n            builder : the system builder\n        Raises:\n            ValueError: Raises an error when trainer_networks is not\n                        set to single_trainer.\n        \"\"\"\n        unique_net_keys = builder.store.unique_net_keys\n\n        # Setup trainer_networks\n\n        builder.store.trainer_networks = {\"trainer\": unique_net_keys}\n\n        # Get all the unique trainer network keys\n        all_trainer_net_keys = []\n        for trainer_nets in builder.store.trainer_networks.values():\n            all_trainer_net_keys.extend(trainer_nets)\n        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))\n\n        # Check that all agent_net_keys are in trainer_networks\n        assert unique_net_keys == unique_trainer_net_keys\n        # Setup specs for each network\n        builder.store.net_spec_keys = {}\n        for i in range(len(unique_net_keys)):\n            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[\n                i % len(builder.store.agents)\n            ]\n\n        # Setup table_network_config\n        builder.store.table_network_config = {}\n        for trainer_key in builder.store.trainer_networks.keys():\n            most_matches = 0\n            trainer_nets = builder.store.trainer_networks[trainer_key]\n            for sample in builder.store.network_sampling_setup:\n                matches = 0\n                for entry in sample:\n                    if entry in trainer_nets:\n                        matches += 1\n                if most_matches < matches:\n                    matches = most_matches\n                    builder.store.table_network_config[trainer_key] = sample\n\n        builder.store.networks = builder.store.network_factory()\n\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"_summary_\"\"\"\n        # Convert network keys for the trainer.\n        trainer.store.trainer_table_entry = trainer.store.table_network_config[\n            trainer.store.trainer_id\n        ]\n        trainer.store.trainer_agents = trainer.store.agents[\n            : len(trainer.store.trainer_table_entry)\n        ]\n        trainer.store.trainer_agent_net_keys = {\n            agent: trainer.store.trainer_table_entry[a_i]\n            for a_i, agent in enumerate(trainer.store.trainer_agents)\n        }\n\n\nclass OneTrainerPerNetworkInit(BaseTrainerInit):\n    def __init__(self, config: SimpleNamespace = SimpleNamespace()):\n        \"\"\"Initialises a multiple trainers.\n\n        Different trainer will be dedicated to training each network.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\".\n\n        Args:\n            builder : the system builder\n        Raises:\n            ValueError: Raises an error when trainer_networks is not\n                        set to one_trainer_per_network.\n        \"\"\"\n        unique_net_keys = builder.store.unique_net_keys\n\n        # Setup trainer_networks\n        builder.store.trainer_networks = {\n            f\"trainer_{i}\": [unique_net_keys[i]] for i in range(len(unique_net_keys))\n        }\n\n        # Get all the unique trainer network keys\n        all_trainer_net_keys = []\n        for trainer_nets in builder.store.trainer_networks.values():\n            all_trainer_net_keys.extend(trainer_nets)\n        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))\n\n        # Check that all agent_net_keys are in trainer_networks\n        assert unique_net_keys == unique_trainer_net_keys\n        # Setup specs for each network\n        builder.store.net_spec_keys = {}\n        for i in range(len(unique_net_keys)):\n            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[\n                i % len(builder.store.agents)\n            ]\n\n        # Setup table_network_config\n        builder.store.table_network_config = {}\n        for trainer_key in builder.store.trainer_networks.keys():\n            most_matches = 0\n            trainer_nets = builder.store.trainer_networks[trainer_key]\n            for sample in builder.store.network_sampling_setup:\n                matches = 0\n                for entry in sample:\n                    if entry in trainer_nets:\n                        matches += 1\n                if most_matches < matches:\n                    matches = most_matches\n                    builder.store.table_network_config[trainer_key] = sample\n\n        builder.store.networks = builder.store.network_factory()\n\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"_summary_\"\"\"\n        # Convert network keys for the trainer.\n        trainer.store.trainer_table_entry = trainer.store.table_network_config[\n            trainer.store.trainer_id\n        ]\n        trainer.store.trainer_agents = trainer.store.agents[\n            : len(trainer.store.trainer_table_entry)\n        ]\n        trainer.store.trainer_agent_net_keys = {\n            agent: trainer.store.trainer_table_entry[a_i]\n            for a_i, agent in enumerate(trainer.store.trainer_agents)\n        }\n\n\n@dataclass\nclass CustomTrainerInitConfig:\n    trainer_networks: Dict = field(default_factory=lambda: {})\n\n\nclass CustomTrainerInit(BaseTrainerInit):\n    def __init__(self, config: CustomTrainerInitConfig = CustomTrainerInitConfig()):\n        \"\"\"Initialises custom trainers.\n\n        Custom trainer network configuration can be given as a dictionary\n        assigning specific trainers to specific networks.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n\n        self.config = config\n\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\"Assigns trainers to networks for training.\n\n        Args:\n            builder : the system builder\n        Raises:\n            ValueError: Raises an error when trainer_networks is not\n                        passed in as a dictionary.\n        \"\"\"\n        trainer_networks = self.config.trainer_networks\n        unique_net_keys = builder.store.unique_net_keys\n\n        # Setup trainer_networks\n        if not isinstance(trainer_networks, dict) or trainer_networks == {}:\n\n            raise ValueError(\"trainer_networks must be a dictionary.\")\n\n        builder.store.trainer_networks = trainer_networks\n\n        # Get all the unique trainer network keys\n        all_trainer_net_keys = []\n        for trainer_nets in builder.store.trainer_networks.values():\n            all_trainer_net_keys.extend(trainer_nets)\n        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))\n\n        # Check that all agent_net_keys are in trainer_networks\n        assert unique_net_keys == unique_trainer_net_keys\n        # Setup specs for each network\n        builder.store.net_spec_keys = {}\n        for i in range(len(unique_net_keys)):\n            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[\n                i % len(builder.store.agents)\n            ]\n\n        # Setup table_network_config\n        builder.store.table_network_config = {}\n        for trainer_key in builder.store.trainer_networks.keys():\n            most_matches = 0\n            trainer_nets = builder.store.trainer_networks[trainer_key]\n            for sample in builder.store.network_sampling_setup:\n                matches = 0\n                for entry in sample:\n                    if entry in trainer_nets:\n                        matches += 1\n                if most_matches < matches:\n                    matches = most_matches\n                    builder.store.table_network_config[trainer_key] = sample\n\n        builder.store.networks = builder.store.network_factory()\n\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"_summary_\"\"\"\n        # Convert network keys for the trainer.\n        trainer.store.trainer_table_entry = trainer.store.table_network_config[\n            trainer.store.trainer_id\n        ]\n        trainer.store.trainer_agents = trainer.store.agents[\n            : len(trainer.store.trainer_table_entry)\n        ]\n        trainer.store.trainer_agent_net_keys = {\n            agent: trainer.store.trainer_table_entry[a_i]\n            for a_i, agent in enumerate(trainer.store.trainer_agents)\n        }\n\n    @staticmethod\n    def config_class() -> Optional[Callable]:\n        \"\"\"Config class used for component.\n\n        Returns:\n            config class/dataclass for component.\n        \"\"\"\n        return CustomTrainerInitConfig\n", "path": "mava/components/jax/training/trainer.py"}], "after_files": [{"content": "# python3\n# Copyright 2021 InstaDeep Ltd. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Trainer components for system builders.\"\"\"\n\nimport abc\nfrom dataclasses import dataclass, field\nfrom types import SimpleNamespace\nfrom typing import Any, Callable, Dict, Optional\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemBuilder, SystemTrainer\nfrom mava.utils.sort_utils import sort_str_num\n\n\nclass BaseTrainerInit(Component):\n    @abc.abstractmethod\n    def __init__(\n        self,\n        config: Any,\n    ):\n        \"\"\"Initialise system init components.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    @abc.abstractmethod\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\"Summary.\"\"\"\n        pass\n\n    @abc.abstractmethod\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"Summary.\"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"Component name.\"\"\"\n\n        return \"trainer_init\"\n\n\nclass SingleTrainerInit(BaseTrainerInit):\n    def __init__(self, config: SimpleNamespace = SimpleNamespace()):\n        \"\"\"Initialises a single trainer.\n\n        Single trainer is used to train all networks.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\"Assigns trainers to networks for training.\n\n        Args:\n            builder : the system builder\n        Raises:\n            ValueError: Raises an error when trainer_networks is not\n                        set to single_trainer.\n        \"\"\"\n        unique_net_keys = builder.store.unique_net_keys\n\n        # Setup trainer_networks\n\n        builder.store.trainer_networks = {\"trainer\": unique_net_keys}\n\n        # Get all the unique trainer network keys\n        all_trainer_net_keys = []\n        for trainer_nets in builder.store.trainer_networks.values():\n            all_trainer_net_keys.extend(trainer_nets)\n        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))\n\n        # Check that all agent_net_keys are in trainer_networks\n        assert unique_net_keys == unique_trainer_net_keys\n        # Setup specs for each network\n        builder.store.net_spec_keys = {}\n        for i in range(len(unique_net_keys)):\n            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[\n                i % len(builder.store.agents)\n            ]\n\n        # Setup table_network_config\n        builder.store.table_network_config = {}\n        for trainer_key in builder.store.trainer_networks.keys():\n            most_matches = 0\n            trainer_nets = builder.store.trainer_networks[trainer_key]\n            for sample in builder.store.network_sampling_setup:\n                matches = 0\n                for entry in sample:\n                    if entry in trainer_nets:\n                        matches += 1\n                if most_matches < matches:\n                    matches = most_matches\n                    builder.store.table_network_config[trainer_key] = sample\n\n        builder.store.networks = builder.store.network_factory()\n\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"_summary_\"\"\"\n        # Convert network keys for the trainer.\n        trainer.store.trainer_table_entry = trainer.store.table_network_config[\n            trainer.store.trainer_id\n        ]\n        trainer.store.trainer_agents = trainer.store.agents[\n            : len(trainer.store.trainer_table_entry)\n        ]\n        trainer.store.trainer_agent_net_keys = {\n            agent: trainer.store.trainer_table_entry[a_i]\n            for a_i, agent in enumerate(trainer.store.trainer_agents)\n        }\n\n\nclass OneTrainerPerNetworkInit(BaseTrainerInit):\n    def __init__(self, config: SimpleNamespace = SimpleNamespace()):\n        \"\"\"Initialises a multiple trainers.\n\n        Different trainer will be dedicated to training each network.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\".\n\n        Args:\n            builder : the system builder\n        Raises:\n            ValueError: Raises an error when trainer_networks is not\n                        set to one_trainer_per_network.\n        \"\"\"\n        unique_net_keys = builder.store.unique_net_keys\n\n        # Setup trainer_networks\n        builder.store.trainer_networks = {\n            f\"trainer_{i}\": [unique_net_keys[i]] for i in range(len(unique_net_keys))\n        }\n\n        # Get all the unique trainer network keys\n        all_trainer_net_keys = []\n        for trainer_nets in builder.store.trainer_networks.values():\n            all_trainer_net_keys.extend(trainer_nets)\n        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))\n\n        # Check that all agent_net_keys are in trainer_networks\n        assert unique_net_keys == unique_trainer_net_keys\n        # Setup specs for each network\n        builder.store.net_spec_keys = {}\n        for i in range(len(unique_net_keys)):\n            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[\n                i % len(builder.store.agents)\n            ]\n\n        # Setup table_network_config\n        builder.store.table_network_config = {}\n        for trainer_key in builder.store.trainer_networks.keys():\n            most_matches = 0\n            trainer_nets = builder.store.trainer_networks[trainer_key]\n            for sample in builder.store.network_sampling_setup:\n                matches = 0\n                for entry in sample:\n                    if entry in trainer_nets:\n                        matches += 1\n                if most_matches < matches:\n                    matches = most_matches\n                    builder.store.table_network_config[trainer_key] = sample\n\n        builder.store.networks = builder.store.network_factory()\n\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"_summary_\"\"\"\n        # Convert network keys for the trainer.\n        trainer.store.trainer_table_entry = trainer.store.table_network_config[\n            trainer.store.trainer_id\n        ]\n        trainer.store.trainer_agents = trainer.store.agents[\n            : len(trainer.store.trainer_table_entry)\n        ]\n        trainer.store.trainer_agent_net_keys = {\n            agent: trainer.store.trainer_table_entry[a_i]\n            for a_i, agent in enumerate(trainer.store.trainer_agents)\n        }\n\n\n@dataclass\nclass CustomTrainerInitConfig:\n    trainer_networks: Dict = field(default_factory=lambda: {})\n\n\nclass CustomTrainerInit(BaseTrainerInit):\n    def __init__(self, config: CustomTrainerInitConfig = CustomTrainerInitConfig()):\n        \"\"\"Initialises custom trainers.\n\n        Custom trainer network configuration can be given as a dictionary\n        assigning specific trainers to specific networks.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n\n        self.config = config\n\n    def on_building_init_end(self, builder: SystemBuilder) -> None:\n        \"\"\"Assigns trainers to networks for training.\n\n        Args:\n            builder : the system builder\n        Raises:\n            ValueError: Raises an error when trainer_networks is not\n                        passed in as a dictionary.\n        \"\"\"\n        trainer_networks = self.config.trainer_networks\n        unique_net_keys = builder.store.unique_net_keys\n\n        # Setup trainer_networks\n        if not isinstance(trainer_networks, dict) or trainer_networks == {}:\n\n            raise ValueError(\"trainer_networks must be a non-empty dictionary.\")\n\n        builder.store.trainer_networks = trainer_networks\n\n        # Get all the unique trainer network keys\n        all_trainer_net_keys = []\n        for trainer_nets in builder.store.trainer_networks.values():\n            all_trainer_net_keys.extend(trainer_nets)\n        unique_trainer_net_keys = sort_str_num(list(set(all_trainer_net_keys)))\n\n        # Check that all agent_net_keys are in trainer_networks\n        assert unique_net_keys == unique_trainer_net_keys\n        # Setup specs for each network\n        builder.store.net_spec_keys = {}\n        for i in range(len(unique_net_keys)):\n            builder.store.net_spec_keys[unique_net_keys[i]] = builder.store.agents[\n                i % len(builder.store.agents)\n            ]\n\n        # Setup table_network_config\n        builder.store.table_network_config = {}\n        for trainer_key in builder.store.trainer_networks.keys():\n            most_matches = 0\n            trainer_nets = builder.store.trainer_networks[trainer_key]\n            for sample in builder.store.network_sampling_setup:\n                matches = 0\n                for entry in sample:\n                    if entry in trainer_nets:\n                        matches += 1\n                if most_matches < matches:\n                    matches = most_matches\n                    builder.store.table_network_config[trainer_key] = sample\n\n        builder.store.networks = builder.store.network_factory()\n\n    def on_training_utility_fns(self, trainer: SystemTrainer) -> None:\n        \"\"\"_summary_\"\"\"\n        # Convert network keys for the trainer.\n        trainer.store.trainer_table_entry = trainer.store.table_network_config[\n            trainer.store.trainer_id\n        ]\n        trainer.store.trainer_agents = trainer.store.agents[\n            : len(trainer.store.trainer_table_entry)\n        ]\n        trainer.store.trainer_agent_net_keys = {\n            agent: trainer.store.trainer_table_entry[a_i]\n            for a_i, agent in enumerate(trainer.store.trainer_agents)\n        }\n\n    @staticmethod\n    def config_class() -> Optional[Callable]:\n        \"\"\"Config class used for component.\n\n        Returns:\n            config class/dataclass for component.\n        \"\"\"\n        return CustomTrainerInitConfig\n", "path": "mava/components/jax/training/trainer.py"}]}