We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[TEST] Jax Building Base
### What do you want to test?
Jax SystemInit component

### 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/building/system_init.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 import abc
17 from dataclasses import dataclass, field
18 from typing import Any, Callable, List, Optional, Union
19 
20 from mava.components.jax import Component
21 from mava.core_jax import SystemBuilder
22 from mava.utils import enums
23 from mava.utils.sort_utils import sample_new_agent_keys, sort_str_num
24 
25 
26 class BaseSystemInit(Component):
27     @abc.abstractmethod
28     def __init__(
29         self,
30         config: Any,
31     ):
32         """Initialise system init components.
33 
34         Args:
35             config : a dataclass specifying the component parameters.
36         """
37         self.config = config
38 
39     @abc.abstractmethod
40     def on_building_init(self, builder: SystemBuilder) -> None:
41         """Compute and add network sampling information to the builder."""
42         pass
43 
44     @staticmethod
45     def name() -> str:
46         """_summary_"""
47         return "system_init"
48 
49 
50 @dataclass
51 class FixedNetworkSystemInitConfig:
52     shared_weights: bool = True
53     network_sampling_setup: enums.NetworkSampler = (
54         enums.NetworkSampler.fixed_agent_networks
55     )
56 
57 
58 class FixedNetworkSystemInit(BaseSystemInit):
59     def __init__(
60         self, config: FixedNetworkSystemInitConfig = FixedNetworkSystemInitConfig()
61     ):
62         """Agents will always be assigned to a fixed network.
63 
64         A single network get assigned to all agents of the same type
65         if weights are shared, else a seperate network is assigned to each agent
66 
67         Args:
68             config : a dataclass specifying the component parameters
69         """
70         self.config = config
71 
72     def on_building_init(self, builder: SystemBuilder) -> None:
73         """Compute and add network sampling information to the builder."""
74         if (
75             self.config.network_sampling_setup
76             != enums.NetworkSampler.fixed_agent_networks
77         ):
78             raise ValueError(
79                 "Fixed network system init requires fixed_agent_networks sampling"
80             )
81 
82         builder.store.agent_net_keys = {
83             agent: f"network_{agent.split('_')[0]}"
84             if self.config.shared_weights
85             else f"network_{agent}"
86             for agent in builder.store.agents
87         }
88         builder.store.network_sampling_setup = [
89             [
90                 builder.store.agent_net_keys[key]
91                 for key in sort_str_num(builder.store.agent_net_keys.keys())
92             ]
93         ]
94 
95         # Check that the environment and agent_net_keys has the same amount of agents
96         sample_length = len(builder.store.network_sampling_setup[0])
97         agent_ids = builder.store.environment_spec.get_agent_ids()
98         assert len(agent_ids) == len(builder.store.agent_net_keys.keys())
99 
100         # Check if the samples are of the same length and that they perfectly fit
101         # into the total number of agents
102         assert len(builder.store.agent_net_keys.keys()) % sample_length == 0
103         for i in range(1, len(builder.store.network_sampling_setup)):
104             assert len(builder.store.network_sampling_setup[i]) == sample_length
105 
106         # Get all the unique agent network keys
107         all_samples = []
108         for sample in builder.store.network_sampling_setup:
109             all_samples.extend(sample)
110         builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))
111 
112         # Create mapping from networks to ints
113         builder.store.net_keys_to_ids = {
114             net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)
115         }
116 
117     @staticmethod
118     def config_class() -> Optional[Callable]:
119         """Config class used for component.
120 
121         Returns:
122             config class/dataclass for component.
123         """
124         return FixedNetworkSystemInitConfig
125 
126 
127 @dataclass
128 class RandomSamplingSystemInitConfig:
129     network_sampling_setup: enums.NetworkSampler = (
130         enums.NetworkSampler.random_agent_networks
131     )
132     shared_weights: bool = False
133 
134 
135 class RandomSamplingSystemInit(BaseSystemInit):
136     def __init__(
137         self, config: RandomSamplingSystemInitConfig = RandomSamplingSystemInitConfig()
138     ):
139         """Create N network policies, where N is the number of agents.
140 
141         Randomly select policies from this set for each agent at the start of a
142         episode. This sampling is done with replacement so the same policy
143         can be selected for more than one agent for a given episode.
144 
145         Args:
146             config : a dataclass specifying the component parameters.
147         """
148         self.config = config
149 
150     def on_building_init(self, builder: SystemBuilder) -> None:
151         """Compute and add network sampling information to the builder."""
152         if (
153             self.config.network_sampling_setup
154             != enums.NetworkSampler.random_agent_networks
155         ):
156             raise ValueError(
157                 "Random sampling system init requires random_agent_networks sampling"
158             )
159 
160         if builder.store.shared_weights:
161             raise ValueError(
162                 "Shared weights cannot be used with random policy per agent"
163             )
164         builder.store.agent_net_keys = {
165             builder.store.agents[i]: f"network_{i}"
166             for i in range(len(builder.store.agents))
167         }
168 
169         builder.store.network_sampling_setup = [
170             [builder.store.agent_net_keys[key]]
171             for key in sort_str_num(builder.store.agent_net_keys.keys())
172         ]
173 
174         # Check that the environment and agent_net_keys has the same amount of agents
175         sample_length = len(builder.store.network_sampling_setup[0])
176         agent_ids = builder.store.environment_spec.get_agent_ids()
177         assert len(agent_ids) == len(builder.store.agent_net_keys.keys())
178 
179         # Check if the samples are of the same length and that they perfectly fit
180         # into the total number of agents
181         assert len(builder.store.agent_net_keys.keys()) % sample_length == 0
182         for i in range(1, len(builder.store.network_sampling_setup)):
183             assert len(builder.store.network_sampling_setup[i]) == sample_length
184 
185         # Get all the unique agent network keys
186         all_samples = []
187         for sample in builder.store.network_sampling_setup:
188             all_samples.extend(sample)
189         builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))
190 
191         # Create mapping from networks to ints
192         builder.store.net_keys_to_ids = {
193             net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)
194         }
195 
196     @staticmethod
197     def config_class() -> Optional[Callable]:
198         """Config class used for component.
199 
200         Returns:
201             config class/dataclass for component.
202         """
203         return RandomSamplingSystemInitConfig
204 
205 
206 @dataclass
207 class CustomSamplingSystemInitConfig:
208     network_sampling_setup: Union[List, enums.NetworkSampler] = field(
209         default_factory=lambda: []
210     )
211     shared_weights: bool = False
212 
213 
214 class CustomSamplingSystemInit(BaseSystemInit):
215     def __init__(
216         self,
217         config: CustomSamplingSystemInitConfig = CustomSamplingSystemInitConfig(),
218     ):
219         """Specify a custom network sampling setup.
220 
221         This network sampling setup will then get used to randomly select
222         policies for agents.
223 
224         Args:
225             config : a dataclass specifying the component parameters.
226         """
227         self.config = config
228 
229     def on_building_init(self, builder: SystemBuilder) -> None:
230         """Compute and add network sampling information to the builder."""
231         if self.config.network_sampling_setup == []:
232             raise ValueError("A custom network sampling setup list must be provided.")
233 
234         # Use network_sampling_setup to determine agent_net_keys
235         _, builder.store.agent_net_keys = sample_new_agent_keys(
236             builder.store.agents,
237             self.config.network_sampling_setup,
238         )
239 
240         # Check that the environment and agent_net_keys has the same amount of agents
241         sample_length = len(self.config.network_sampling_setup[0])
242         agent_ids = builder.store.environment_spec.get_agent_ids()
243         assert len(agent_ids) == len(builder.store.agent_net_keys.keys())
244 
245         # Check if the samples are of the same length and that they perfectly fit
246         # into the total number of agents
247         assert len(builder.store.agent_net_keys.keys()) % sample_length == 0
248         for i in range(1, len(self.config.network_sampling_setup)):
249             assert len(self.config.network_sampling_setup[i]) == sample_length
250 
251         # Get all the unique agent network keys
252         all_samples = []
253         for sample in self.config.network_sampling_setup:
254             all_samples.extend(sample)
255         builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))
256 
257         # Create mapping from networks to ints
258         builder.store.net_keys_to_ids = {
259             net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)
260         }
261 
262     @staticmethod
263     def config_class() -> Optional[Callable]:
264         """Config class used for component.
265 
266         Returns:
267             config class/dataclass for component.
268         """
269         return CustomSamplingSystemInitConfig
270
```
--- 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/building/system_init.py b/mava/components/jax/building/system_init.py
--- a/mava/components/jax/building/system_init.py
+++ b/mava/components/jax/building/system_init.py
@@ -48,8 +48,12 @@
 
 
 @dataclass
-class FixedNetworkSystemInitConfig:
+class BaseSystemInitConfig:
     shared_weights: bool = True
+
+
+@dataclass
+class FixedNetworkSystemInitConfig(BaseSystemInitConfig):
     network_sampling_setup: enums.NetworkSampler = (
         enums.NetworkSampler.fixed_agent_networks
     )
@@ -129,7 +133,6 @@
     network_sampling_setup: enums.NetworkSampler = (
         enums.NetworkSampler.random_agent_networks
     )
-    shared_weights: bool = False
 
 
 class RandomSamplingSystemInit(BaseSystemInit):
@@ -157,10 +160,6 @@
                 "Random sampling system init requires random_agent_networks sampling"
             )
 
-        if builder.store.shared_weights:
-            raise ValueError(
-                "Shared weights cannot be used with random policy per agent"
-            )
         builder.store.agent_net_keys = {
             builder.store.agents[i]: f"network_{i}"
             for i in range(len(builder.store.agents))
@@ -204,11 +203,10 @@
 
 
 @dataclass
-class CustomSamplingSystemInitConfig:
+class CustomSamplingSystemInitConfig(BaseSystemInitConfig):
     network_sampling_setup: Union[List, enums.NetworkSampler] = field(
         default_factory=lambda: []
     )
-    shared_weights: bool = False
 
 
 class CustomSamplingSystemInit(BaseSystemInit):

{"golden_diff": "diff --git a/mava/components/jax/building/system_init.py b/mava/components/jax/building/system_init.py\n--- a/mava/components/jax/building/system_init.py\n+++ b/mava/components/jax/building/system_init.py\n@@ -48,8 +48,12 @@\n \n \n @dataclass\n-class FixedNetworkSystemInitConfig:\n+class BaseSystemInitConfig:\n     shared_weights: bool = True\n+\n+\n+@dataclass\n+class FixedNetworkSystemInitConfig(BaseSystemInitConfig):\n     network_sampling_setup: enums.NetworkSampler = (\n         enums.NetworkSampler.fixed_agent_networks\n     )\n@@ -129,7 +133,6 @@\n     network_sampling_setup: enums.NetworkSampler = (\n         enums.NetworkSampler.random_agent_networks\n     )\n-    shared_weights: bool = False\n \n \n class RandomSamplingSystemInit(BaseSystemInit):\n@@ -157,10 +160,6 @@\n                 \"Random sampling system init requires random_agent_networks sampling\"\n             )\n \n-        if builder.store.shared_weights:\n-            raise ValueError(\n-                \"Shared weights cannot be used with random policy per agent\"\n-            )\n         builder.store.agent_net_keys = {\n             builder.store.agents[i]: f\"network_{i}\"\n             for i in range(len(builder.store.agents))\n@@ -204,11 +203,10 @@\n \n \n @dataclass\n-class CustomSamplingSystemInitConfig:\n+class CustomSamplingSystemInitConfig(BaseSystemInitConfig):\n     network_sampling_setup: Union[List, enums.NetworkSampler] = field(\n         default_factory=lambda: []\n     )\n-    shared_weights: bool = False\n \n \n class CustomSamplingSystemInit(BaseSystemInit):\n", "issue": "[TEST] Jax Building Base\n### What do you want to test?\r\nJax SystemInit component\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\nimport abc\nfrom dataclasses import dataclass, field\nfrom typing import Any, Callable, List, Optional, Union\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemBuilder\nfrom mava.utils import enums\nfrom mava.utils.sort_utils import sample_new_agent_keys, sort_str_num\n\n\nclass BaseSystemInit(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(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\"\"\"\n        return \"system_init\"\n\n\n@dataclass\nclass FixedNetworkSystemInitConfig:\n    shared_weights: bool = True\n    network_sampling_setup: enums.NetworkSampler = (\n        enums.NetworkSampler.fixed_agent_networks\n    )\n\n\nclass FixedNetworkSystemInit(BaseSystemInit):\n    def __init__(\n        self, config: FixedNetworkSystemInitConfig = FixedNetworkSystemInitConfig()\n    ):\n        \"\"\"Agents will always be assigned to a fixed network.\n\n        A single network get assigned to all agents of the same type\n        if weights are shared, else a seperate network is assigned to each agent\n\n        Args:\n            config : a dataclass specifying the component parameters\n        \"\"\"\n        self.config = config\n\n    def on_building_init(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        if (\n            self.config.network_sampling_setup\n            != enums.NetworkSampler.fixed_agent_networks\n        ):\n            raise ValueError(\n                \"Fixed network system init requires fixed_agent_networks sampling\"\n            )\n\n        builder.store.agent_net_keys = {\n            agent: f\"network_{agent.split('_')[0]}\"\n            if self.config.shared_weights\n            else f\"network_{agent}\"\n            for agent in builder.store.agents\n        }\n        builder.store.network_sampling_setup = [\n            [\n                builder.store.agent_net_keys[key]\n                for key in sort_str_num(builder.store.agent_net_keys.keys())\n            ]\n        ]\n\n        # Check that the environment and agent_net_keys has the same amount of agents\n        sample_length = len(builder.store.network_sampling_setup[0])\n        agent_ids = builder.store.environment_spec.get_agent_ids()\n        assert len(agent_ids) == len(builder.store.agent_net_keys.keys())\n\n        # Check if the samples are of the same length and that they perfectly fit\n        # into the total number of agents\n        assert len(builder.store.agent_net_keys.keys()) % sample_length == 0\n        for i in range(1, len(builder.store.network_sampling_setup)):\n            assert len(builder.store.network_sampling_setup[i]) == sample_length\n\n        # Get all the unique agent network keys\n        all_samples = []\n        for sample in builder.store.network_sampling_setup:\n            all_samples.extend(sample)\n        builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))\n\n        # Create mapping from networks to ints\n        builder.store.net_keys_to_ids = {\n            net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)\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 FixedNetworkSystemInitConfig\n\n\n@dataclass\nclass RandomSamplingSystemInitConfig:\n    network_sampling_setup: enums.NetworkSampler = (\n        enums.NetworkSampler.random_agent_networks\n    )\n    shared_weights: bool = False\n\n\nclass RandomSamplingSystemInit(BaseSystemInit):\n    def __init__(\n        self, config: RandomSamplingSystemInitConfig = RandomSamplingSystemInitConfig()\n    ):\n        \"\"\"Create N network policies, where N is the number of agents.\n\n        Randomly select policies from this set for each agent at the start of a\n        episode. This sampling is done with replacement so the same policy\n        can be selected for more than one agent for a given episode.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        if (\n            self.config.network_sampling_setup\n            != enums.NetworkSampler.random_agent_networks\n        ):\n            raise ValueError(\n                \"Random sampling system init requires random_agent_networks sampling\"\n            )\n\n        if builder.store.shared_weights:\n            raise ValueError(\n                \"Shared weights cannot be used with random policy per agent\"\n            )\n        builder.store.agent_net_keys = {\n            builder.store.agents[i]: f\"network_{i}\"\n            for i in range(len(builder.store.agents))\n        }\n\n        builder.store.network_sampling_setup = [\n            [builder.store.agent_net_keys[key]]\n            for key in sort_str_num(builder.store.agent_net_keys.keys())\n        ]\n\n        # Check that the environment and agent_net_keys has the same amount of agents\n        sample_length = len(builder.store.network_sampling_setup[0])\n        agent_ids = builder.store.environment_spec.get_agent_ids()\n        assert len(agent_ids) == len(builder.store.agent_net_keys.keys())\n\n        # Check if the samples are of the same length and that they perfectly fit\n        # into the total number of agents\n        assert len(builder.store.agent_net_keys.keys()) % sample_length == 0\n        for i in range(1, len(builder.store.network_sampling_setup)):\n            assert len(builder.store.network_sampling_setup[i]) == sample_length\n\n        # Get all the unique agent network keys\n        all_samples = []\n        for sample in builder.store.network_sampling_setup:\n            all_samples.extend(sample)\n        builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))\n\n        # Create mapping from networks to ints\n        builder.store.net_keys_to_ids = {\n            net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)\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 RandomSamplingSystemInitConfig\n\n\n@dataclass\nclass CustomSamplingSystemInitConfig:\n    network_sampling_setup: Union[List, enums.NetworkSampler] = field(\n        default_factory=lambda: []\n    )\n    shared_weights: bool = False\n\n\nclass CustomSamplingSystemInit(BaseSystemInit):\n    def __init__(\n        self,\n        config: CustomSamplingSystemInitConfig = CustomSamplingSystemInitConfig(),\n    ):\n        \"\"\"Specify a custom network sampling setup.\n\n        This network sampling setup will then get used to randomly select\n        policies for agents.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        if self.config.network_sampling_setup == []:\n            raise ValueError(\"A custom network sampling setup list must be provided.\")\n\n        # Use network_sampling_setup to determine agent_net_keys\n        _, builder.store.agent_net_keys = sample_new_agent_keys(\n            builder.store.agents,\n            self.config.network_sampling_setup,\n        )\n\n        # Check that the environment and agent_net_keys has the same amount of agents\n        sample_length = len(self.config.network_sampling_setup[0])\n        agent_ids = builder.store.environment_spec.get_agent_ids()\n        assert len(agent_ids) == len(builder.store.agent_net_keys.keys())\n\n        # Check if the samples are of the same length and that they perfectly fit\n        # into the total number of agents\n        assert len(builder.store.agent_net_keys.keys()) % sample_length == 0\n        for i in range(1, len(self.config.network_sampling_setup)):\n            assert len(self.config.network_sampling_setup[i]) == sample_length\n\n        # Get all the unique agent network keys\n        all_samples = []\n        for sample in self.config.network_sampling_setup:\n            all_samples.extend(sample)\n        builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))\n\n        # Create mapping from networks to ints\n        builder.store.net_keys_to_ids = {\n            net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)\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 CustomSamplingSystemInitConfig\n", "path": "mava/components/jax/building/system_init.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\nimport abc\nfrom dataclasses import dataclass, field\nfrom typing import Any, Callable, List, Optional, Union\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemBuilder\nfrom mava.utils import enums\nfrom mava.utils.sort_utils import sample_new_agent_keys, sort_str_num\n\n\nclass BaseSystemInit(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(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\"\"\"\n        return \"system_init\"\n\n\n@dataclass\nclass BaseSystemInitConfig:\n    shared_weights: bool = True\n\n\n@dataclass\nclass FixedNetworkSystemInitConfig(BaseSystemInitConfig):\n    network_sampling_setup: enums.NetworkSampler = (\n        enums.NetworkSampler.fixed_agent_networks\n    )\n\n\nclass FixedNetworkSystemInit(BaseSystemInit):\n    def __init__(\n        self, config: FixedNetworkSystemInitConfig = FixedNetworkSystemInitConfig()\n    ):\n        \"\"\"Agents will always be assigned to a fixed network.\n\n        A single network get assigned to all agents of the same type\n        if weights are shared, else a seperate network is assigned to each agent\n\n        Args:\n            config : a dataclass specifying the component parameters\n        \"\"\"\n        self.config = config\n\n    def on_building_init(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        if (\n            self.config.network_sampling_setup\n            != enums.NetworkSampler.fixed_agent_networks\n        ):\n            raise ValueError(\n                \"Fixed network system init requires fixed_agent_networks sampling\"\n            )\n\n        builder.store.agent_net_keys = {\n            agent: f\"network_{agent.split('_')[0]}\"\n            if self.config.shared_weights\n            else f\"network_{agent}\"\n            for agent in builder.store.agents\n        }\n        builder.store.network_sampling_setup = [\n            [\n                builder.store.agent_net_keys[key]\n                for key in sort_str_num(builder.store.agent_net_keys.keys())\n            ]\n        ]\n\n        # Check that the environment and agent_net_keys has the same amount of agents\n        sample_length = len(builder.store.network_sampling_setup[0])\n        agent_ids = builder.store.environment_spec.get_agent_ids()\n        assert len(agent_ids) == len(builder.store.agent_net_keys.keys())\n\n        # Check if the samples are of the same length and that they perfectly fit\n        # into the total number of agents\n        assert len(builder.store.agent_net_keys.keys()) % sample_length == 0\n        for i in range(1, len(builder.store.network_sampling_setup)):\n            assert len(builder.store.network_sampling_setup[i]) == sample_length\n\n        # Get all the unique agent network keys\n        all_samples = []\n        for sample in builder.store.network_sampling_setup:\n            all_samples.extend(sample)\n        builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))\n\n        # Create mapping from networks to ints\n        builder.store.net_keys_to_ids = {\n            net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)\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 FixedNetworkSystemInitConfig\n\n\n@dataclass\nclass RandomSamplingSystemInitConfig:\n    network_sampling_setup: enums.NetworkSampler = (\n        enums.NetworkSampler.random_agent_networks\n    )\n\n\nclass RandomSamplingSystemInit(BaseSystemInit):\n    def __init__(\n        self, config: RandomSamplingSystemInitConfig = RandomSamplingSystemInitConfig()\n    ):\n        \"\"\"Create N network policies, where N is the number of agents.\n\n        Randomly select policies from this set for each agent at the start of a\n        episode. This sampling is done with replacement so the same policy\n        can be selected for more than one agent for a given episode.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        if (\n            self.config.network_sampling_setup\n            != enums.NetworkSampler.random_agent_networks\n        ):\n            raise ValueError(\n                \"Random sampling system init requires random_agent_networks sampling\"\n            )\n\n        builder.store.agent_net_keys = {\n            builder.store.agents[i]: f\"network_{i}\"\n            for i in range(len(builder.store.agents))\n        }\n\n        builder.store.network_sampling_setup = [\n            [builder.store.agent_net_keys[key]]\n            for key in sort_str_num(builder.store.agent_net_keys.keys())\n        ]\n\n        # Check that the environment and agent_net_keys has the same amount of agents\n        sample_length = len(builder.store.network_sampling_setup[0])\n        agent_ids = builder.store.environment_spec.get_agent_ids()\n        assert len(agent_ids) == len(builder.store.agent_net_keys.keys())\n\n        # Check if the samples are of the same length and that they perfectly fit\n        # into the total number of agents\n        assert len(builder.store.agent_net_keys.keys()) % sample_length == 0\n        for i in range(1, len(builder.store.network_sampling_setup)):\n            assert len(builder.store.network_sampling_setup[i]) == sample_length\n\n        # Get all the unique agent network keys\n        all_samples = []\n        for sample in builder.store.network_sampling_setup:\n            all_samples.extend(sample)\n        builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))\n\n        # Create mapping from networks to ints\n        builder.store.net_keys_to_ids = {\n            net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)\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 RandomSamplingSystemInitConfig\n\n\n@dataclass\nclass CustomSamplingSystemInitConfig(BaseSystemInitConfig):\n    network_sampling_setup: Union[List, enums.NetworkSampler] = field(\n        default_factory=lambda: []\n    )\n\n\nclass CustomSamplingSystemInit(BaseSystemInit):\n    def __init__(\n        self,\n        config: CustomSamplingSystemInitConfig = CustomSamplingSystemInitConfig(),\n    ):\n        \"\"\"Specify a custom network sampling setup.\n\n        This network sampling setup will then get used to randomly select\n        policies for agents.\n\n        Args:\n            config : a dataclass specifying the component parameters.\n        \"\"\"\n        self.config = config\n\n    def on_building_init(self, builder: SystemBuilder) -> None:\n        \"\"\"Compute and add network sampling information to the builder.\"\"\"\n        if self.config.network_sampling_setup == []:\n            raise ValueError(\"A custom network sampling setup list must be provided.\")\n\n        # Use network_sampling_setup to determine agent_net_keys\n        _, builder.store.agent_net_keys = sample_new_agent_keys(\n            builder.store.agents,\n            self.config.network_sampling_setup,\n        )\n\n        # Check that the environment and agent_net_keys has the same amount of agents\n        sample_length = len(self.config.network_sampling_setup[0])\n        agent_ids = builder.store.environment_spec.get_agent_ids()\n        assert len(agent_ids) == len(builder.store.agent_net_keys.keys())\n\n        # Check if the samples are of the same length and that they perfectly fit\n        # into the total number of agents\n        assert len(builder.store.agent_net_keys.keys()) % sample_length == 0\n        for i in range(1, len(self.config.network_sampling_setup)):\n            assert len(self.config.network_sampling_setup[i]) == sample_length\n\n        # Get all the unique agent network keys\n        all_samples = []\n        for sample in self.config.network_sampling_setup:\n            all_samples.extend(sample)\n        builder.store.unique_net_keys = list(sort_str_num(list(set(all_samples))))\n\n        # Create mapping from networks to ints\n        builder.store.net_keys_to_ids = {\n            net_key: i for i, net_key in enumerate(builder.store.unique_net_keys)\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 CustomSamplingSystemInitConfig\n", "path": "mava/components/jax/building/system_init.py"}]}