We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[TEST] Jax Parameter Server
### What do you want to test?
Jax parameter server 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/updating/parameter_server.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 """Parameter server Component for Mava systems."""
17 import abc
18 import time
19 from dataclasses import dataclass
20 from typing import Callable, Optional
21 
22 import numpy as np
23 from acme.jax import savers
24 
25 from mava.components.jax.component import Component
26 from mava.core_jax import SystemParameterServer
27 
28 
29 @dataclass
30 class ParameterServerConfig:
31     checkpoint: bool = True
32     checkpoint_subpath: str = "~/mava/"
33     checkpoint_minute_interval: int = 5
34     non_blocking_sleep_seconds: int = 10
35 
36 
37 class ParameterServer(Component):
38     @abc.abstractmethod
39     def __init__(
40         self,
41         config: ParameterServerConfig = ParameterServerConfig(),
42     ) -> None:
43         """Mock system Component."""
44         self.config = config
45 
46     @abc.abstractmethod
47     def on_parameter_server_init_start(self, server: SystemParameterServer) -> None:
48         """_summary_
49 
50         Args:
51             server : _description_
52         """
53         pass
54 
55     # Get
56     @abc.abstractmethod
57     def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:
58         """_summary_"""
59         pass
60 
61     # Set
62     @abc.abstractmethod
63     def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:
64         """_summary_"""
65         pass
66 
67     # Add
68     @abc.abstractmethod
69     def on_parameter_server_add_to_parameters(
70         self, server: SystemParameterServer
71     ) -> None:
72         """_summary_"""
73         pass
74 
75     # Save variables using checkpointer
76     @abc.abstractmethod
77     def on_parameter_server_run_loop(self, server: SystemParameterServer) -> None:
78         """_summary_
79 
80         Args:
81             server : _description_
82         """
83         pass
84 
85     @staticmethod
86     def name() -> str:
87         """Component type name, e.g. 'dataset' or 'executor'."""
88         return "parameter_server"
89 
90     @staticmethod
91     def config_class() -> Optional[Callable]:
92         """Config class used for Component.
93 
94         Returns:
95             config class/dataclass for Component.
96         """
97         return ParameterServerConfig
98 
99 
100 class DefaultParameterServer(ParameterServer):
101     def __init__(
102         self,
103         config: ParameterServerConfig = ParameterServerConfig(),
104     ) -> None:
105         """Mock system Component."""
106         self.config = config
107 
108     def on_parameter_server_init_start(self, server: SystemParameterServer) -> None:
109         """_summary_
110 
111         Args:
112             server : _description_
113         """
114         networks = server.store.network_factory()
115 
116         # # Create parameters
117         server.store.parameters = {
118             "trainer_steps": np.zeros(1, dtype=np.int32),
119             "trainer_walltime": np.zeros(1, dtype=np.float32),
120             "evaluator_steps": np.zeros(1, dtype=np.int32),
121             "evaluator_episodes": np.zeros(1, dtype=np.int32),
122             "executor_episodes": np.zeros(1, dtype=np.int32),
123             "executor_steps": np.zeros(1, dtype=np.int32),
124         }
125 
126         # Network parameters
127         for net_type_key in networks.keys():
128             for agent_net_key in networks[net_type_key].keys():
129                 # Ensure obs and target networks are sonnet modules
130                 server.store.parameters[f"{net_type_key}-{agent_net_key}"] = networks[
131                     net_type_key
132                 ][agent_net_key].params
133 
134         # Create the checkpointer
135         if self.config.checkpoint:
136             server.store.last_checkpoint_time = 0
137             server.store.checkpoint_minute_interval = (
138                 self.config.checkpoint_minute_interval
139             )
140 
141             # Only save variables that are not empty.
142             save_variables = {}
143             for key in server.store.parameters.keys():
144                 var = server.store.parameters[key]
145                 # Don't store empty tuple (e.g. empty observation_network) variables
146                 if not (type(var) == tuple and len(var) == 0):
147                     save_variables[key] = var
148             server.store.system_checkpointer = savers.Checkpointer(
149                 save_variables, self.config.checkpoint_subpath, time_delta_minutes=0
150             )
151 
152     # Get
153     def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:
154         """_summary_"""
155         names = server.store._param_names
156 
157         if type(names) == str:
158             get_params = server.store.parameters[names]  # type: ignore
159         else:
160             get_params = {}
161             for var_key in names:
162                 get_params[var_key] = server.store.parameters[var_key]
163         server.store.get_parameters = get_params
164 
165     # Set
166     def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:
167         """_summary_"""
168         params = server.store._set_params
169         names = params.keys()
170 
171         if type(names) == str:
172             params = {names: params}  # type: ignore
173             names = [names]  # type: ignore
174 
175         for var_key in names:
176             assert var_key in server.store.parameters
177             if type(server.store.parameters[var_key]) == tuple:
178                 raise NotImplementedError
179                 # # Loop through tuple
180                 # for var_i in range(len(server.store.parameters[var_key])):
181                 #     server.store.parameters[var_key][var_i].assign(params[var_key][var_i])
182             else:
183                 server.store.parameters[var_key] = params[var_key]
184 
185     # Add
186     def on_parameter_server_add_to_parameters(
187         self, server: SystemParameterServer
188     ) -> None:
189         """_summary_"""
190         params = server.store._add_to_params
191         names = params.keys()
192 
193         if type(names) == str:
194             params = {names: params}  # type: ignore
195             names = [names]  # type: ignore
196 
197         for var_key in names:
198             assert var_key in server.store.parameters
199             server.store.parameters[var_key] += params[var_key]
200 
201     # Save variables using checkpointer
202     def on_parameter_server_run_loop(self, server: SystemParameterServer) -> None:
203         """_summary_
204 
205         Args:
206             server : _description_
207         """
208         if (
209             server.store.system_checkpointer
210             and server.store.last_checkpoint_time
211             + server.store.checkpoint_minute_interval * 60
212             + 1
213             < time.time()
214         ):
215             server.store.system_checkpointer.save()
216             server.store.last_checkpoint_time = time.time()
217             print("Updated variables checkpoint.")
218
```
--- 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/updating/parameter_server.py b/mava/components/jax/updating/parameter_server.py
--- a/mava/components/jax/updating/parameter_server.py
+++ b/mava/components/jax/updating/parameter_server.py
@@ -17,7 +17,7 @@
 import abc
 import time
 from dataclasses import dataclass
-from typing import Callable, Optional
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 import numpy as np
 from acme.jax import savers
@@ -134,9 +134,6 @@
         # Create the checkpointer
         if self.config.checkpoint:
             server.store.last_checkpoint_time = 0
-            server.store.checkpoint_minute_interval = (
-                self.config.checkpoint_minute_interval
-            )
 
             # Only save variables that are not empty.
             save_variables = {}
@@ -152,7 +149,7 @@
     # Get
     def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:
         """_summary_"""
-        names = server.store._param_names
+        names: Union[str, Sequence[str]] = server.store._param_names
 
         if type(names) == str:
             get_params = server.store.parameters[names]  # type: ignore
@@ -165,13 +162,9 @@
     # Set
     def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:
         """_summary_"""
-        params = server.store._set_params
+        params: Dict[str, Any] = server.store._set_params
         names = params.keys()
 
-        if type(names) == str:
-            params = {names: params}  # type: ignore
-            names = [names]  # type: ignore
-
         for var_key in names:
             assert var_key in server.store.parameters
             if type(server.store.parameters[var_key]) == tuple:
@@ -187,13 +180,9 @@
         self, server: SystemParameterServer
     ) -> None:
         """_summary_"""
-        params = server.store._add_to_params
+        params: Dict[str, Any] = server.store._add_to_params
         names = params.keys()
 
-        if type(names) == str:
-            params = {names: params}  # type: ignore
-            names = [names]  # type: ignore
-
         for var_key in names:
             assert var_key in server.store.parameters
             server.store.parameters[var_key] += params[var_key]
@@ -206,9 +195,9 @@
             server : _description_
         """
         if (
-            server.store.system_checkpointer
+            self.config.checkpoint
             and server.store.last_checkpoint_time
-            + server.store.checkpoint_minute_interval * 60
+            + self.config.checkpoint_minute_interval * 60
             + 1
             < time.time()
         ):

{"golden_diff": "diff --git a/mava/components/jax/updating/parameter_server.py b/mava/components/jax/updating/parameter_server.py\n--- a/mava/components/jax/updating/parameter_server.py\n+++ b/mava/components/jax/updating/parameter_server.py\n@@ -17,7 +17,7 @@\n import abc\n import time\n from dataclasses import dataclass\n-from typing import Callable, Optional\n+from typing import Any, Callable, Dict, Optional, Sequence, Union\n \n import numpy as np\n from acme.jax import savers\n@@ -134,9 +134,6 @@\n         # Create the checkpointer\n         if self.config.checkpoint:\n             server.store.last_checkpoint_time = 0\n-            server.store.checkpoint_minute_interval = (\n-                self.config.checkpoint_minute_interval\n-            )\n \n             # Only save variables that are not empty.\n             save_variables = {}\n@@ -152,7 +149,7 @@\n     # Get\n     def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:\n         \"\"\"_summary_\"\"\"\n-        names = server.store._param_names\n+        names: Union[str, Sequence[str]] = server.store._param_names\n \n         if type(names) == str:\n             get_params = server.store.parameters[names]  # type: ignore\n@@ -165,13 +162,9 @@\n     # Set\n     def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:\n         \"\"\"_summary_\"\"\"\n-        params = server.store._set_params\n+        params: Dict[str, Any] = server.store._set_params\n         names = params.keys()\n \n-        if type(names) == str:\n-            params = {names: params}  # type: ignore\n-            names = [names]  # type: ignore\n-\n         for var_key in names:\n             assert var_key in server.store.parameters\n             if type(server.store.parameters[var_key]) == tuple:\n@@ -187,13 +180,9 @@\n         self, server: SystemParameterServer\n     ) -> None:\n         \"\"\"_summary_\"\"\"\n-        params = server.store._add_to_params\n+        params: Dict[str, Any] = server.store._add_to_params\n         names = params.keys()\n \n-        if type(names) == str:\n-            params = {names: params}  # type: ignore\n-            names = [names]  # type: ignore\n-\n         for var_key in names:\n             assert var_key in server.store.parameters\n             server.store.parameters[var_key] += params[var_key]\n@@ -206,9 +195,9 @@\n             server : _description_\n         \"\"\"\n         if (\n-            server.store.system_checkpointer\n+            self.config.checkpoint\n             and server.store.last_checkpoint_time\n-            + server.store.checkpoint_minute_interval * 60\n+            + self.config.checkpoint_minute_interval * 60\n             + 1\n             < time.time()\n         ):\n", "issue": "[TEST] Jax Parameter Server\n### What do you want to test?\r\nJax parameter server 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\"\"\"Parameter server Component for Mava systems.\"\"\"\nimport abc\nimport time\nfrom dataclasses import dataclass\nfrom typing import Callable, Optional\n\nimport numpy as np\nfrom acme.jax import savers\n\nfrom mava.components.jax.component import Component\nfrom mava.core_jax import SystemParameterServer\n\n\n@dataclass\nclass ParameterServerConfig:\n    checkpoint: bool = True\n    checkpoint_subpath: str = \"~/mava/\"\n    checkpoint_minute_interval: int = 5\n    non_blocking_sleep_seconds: int = 10\n\n\nclass ParameterServer(Component):\n    @abc.abstractmethod\n    def __init__(\n        self,\n        config: ParameterServerConfig = ParameterServerConfig(),\n    ) -> None:\n        \"\"\"Mock system Component.\"\"\"\n        self.config = config\n\n    @abc.abstractmethod\n    def on_parameter_server_init_start(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        pass\n\n    # Get\n    @abc.abstractmethod\n    def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        pass\n\n    # Set\n    @abc.abstractmethod\n    def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        pass\n\n    # Add\n    @abc.abstractmethod\n    def on_parameter_server_add_to_parameters(\n        self, server: SystemParameterServer\n    ) -> None:\n        \"\"\"_summary_\"\"\"\n        pass\n\n    # Save variables using checkpointer\n    @abc.abstractmethod\n    def on_parameter_server_run_loop(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"Component type name, e.g. 'dataset' or 'executor'.\"\"\"\n        return \"parameter_server\"\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 ParameterServerConfig\n\n\nclass DefaultParameterServer(ParameterServer):\n    def __init__(\n        self,\n        config: ParameterServerConfig = ParameterServerConfig(),\n    ) -> None:\n        \"\"\"Mock system Component.\"\"\"\n        self.config = config\n\n    def on_parameter_server_init_start(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        networks = server.store.network_factory()\n\n        # # Create parameters\n        server.store.parameters = {\n            \"trainer_steps\": np.zeros(1, dtype=np.int32),\n            \"trainer_walltime\": np.zeros(1, dtype=np.float32),\n            \"evaluator_steps\": np.zeros(1, dtype=np.int32),\n            \"evaluator_episodes\": np.zeros(1, dtype=np.int32),\n            \"executor_episodes\": np.zeros(1, dtype=np.int32),\n            \"executor_steps\": np.zeros(1, dtype=np.int32),\n        }\n\n        # Network parameters\n        for net_type_key in networks.keys():\n            for agent_net_key in networks[net_type_key].keys():\n                # Ensure obs and target networks are sonnet modules\n                server.store.parameters[f\"{net_type_key}-{agent_net_key}\"] = networks[\n                    net_type_key\n                ][agent_net_key].params\n\n        # Create the checkpointer\n        if self.config.checkpoint:\n            server.store.last_checkpoint_time = 0\n            server.store.checkpoint_minute_interval = (\n                self.config.checkpoint_minute_interval\n            )\n\n            # Only save variables that are not empty.\n            save_variables = {}\n            for key in server.store.parameters.keys():\n                var = server.store.parameters[key]\n                # Don't store empty tuple (e.g. empty observation_network) variables\n                if not (type(var) == tuple and len(var) == 0):\n                    save_variables[key] = var\n            server.store.system_checkpointer = savers.Checkpointer(\n                save_variables, self.config.checkpoint_subpath, time_delta_minutes=0\n            )\n\n    # Get\n    def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        names = server.store._param_names\n\n        if type(names) == str:\n            get_params = server.store.parameters[names]  # type: ignore\n        else:\n            get_params = {}\n            for var_key in names:\n                get_params[var_key] = server.store.parameters[var_key]\n        server.store.get_parameters = get_params\n\n    # Set\n    def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        params = server.store._set_params\n        names = params.keys()\n\n        if type(names) == str:\n            params = {names: params}  # type: ignore\n            names = [names]  # type: ignore\n\n        for var_key in names:\n            assert var_key in server.store.parameters\n            if type(server.store.parameters[var_key]) == tuple:\n                raise NotImplementedError\n                # # Loop through tuple\n                # for var_i in range(len(server.store.parameters[var_key])):\n                #     server.store.parameters[var_key][var_i].assign(params[var_key][var_i])\n            else:\n                server.store.parameters[var_key] = params[var_key]\n\n    # Add\n    def on_parameter_server_add_to_parameters(\n        self, server: SystemParameterServer\n    ) -> None:\n        \"\"\"_summary_\"\"\"\n        params = server.store._add_to_params\n        names = params.keys()\n\n        if type(names) == str:\n            params = {names: params}  # type: ignore\n            names = [names]  # type: ignore\n\n        for var_key in names:\n            assert var_key in server.store.parameters\n            server.store.parameters[var_key] += params[var_key]\n\n    # Save variables using checkpointer\n    def on_parameter_server_run_loop(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        if (\n            server.store.system_checkpointer\n            and server.store.last_checkpoint_time\n            + server.store.checkpoint_minute_interval * 60\n            + 1\n            < time.time()\n        ):\n            server.store.system_checkpointer.save()\n            server.store.last_checkpoint_time = time.time()\n            print(\"Updated variables checkpoint.\")\n", "path": "mava/components/jax/updating/parameter_server.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\"\"\"Parameter server Component for Mava systems.\"\"\"\nimport abc\nimport time\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Dict, Optional, Sequence, Union\n\nimport numpy as np\nfrom acme.jax import savers\n\nfrom mava.components.jax.component import Component\nfrom mava.core_jax import SystemParameterServer\n\n\n@dataclass\nclass ParameterServerConfig:\n    checkpoint: bool = True\n    checkpoint_subpath: str = \"~/mava/\"\n    checkpoint_minute_interval: int = 5\n    non_blocking_sleep_seconds: int = 10\n\n\nclass ParameterServer(Component):\n    @abc.abstractmethod\n    def __init__(\n        self,\n        config: ParameterServerConfig = ParameterServerConfig(),\n    ) -> None:\n        \"\"\"Mock system Component.\"\"\"\n        self.config = config\n\n    @abc.abstractmethod\n    def on_parameter_server_init_start(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        pass\n\n    # Get\n    @abc.abstractmethod\n    def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        pass\n\n    # Set\n    @abc.abstractmethod\n    def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        pass\n\n    # Add\n    @abc.abstractmethod\n    def on_parameter_server_add_to_parameters(\n        self, server: SystemParameterServer\n    ) -> None:\n        \"\"\"_summary_\"\"\"\n        pass\n\n    # Save variables using checkpointer\n    @abc.abstractmethod\n    def on_parameter_server_run_loop(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"Component type name, e.g. 'dataset' or 'executor'.\"\"\"\n        return \"parameter_server\"\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 ParameterServerConfig\n\n\nclass DefaultParameterServer(ParameterServer):\n    def __init__(\n        self,\n        config: ParameterServerConfig = ParameterServerConfig(),\n    ) -> None:\n        \"\"\"Mock system Component.\"\"\"\n        self.config = config\n\n    def on_parameter_server_init_start(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        networks = server.store.network_factory()\n\n        # # Create parameters\n        server.store.parameters = {\n            \"trainer_steps\": np.zeros(1, dtype=np.int32),\n            \"trainer_walltime\": np.zeros(1, dtype=np.float32),\n            \"evaluator_steps\": np.zeros(1, dtype=np.int32),\n            \"evaluator_episodes\": np.zeros(1, dtype=np.int32),\n            \"executor_episodes\": np.zeros(1, dtype=np.int32),\n            \"executor_steps\": np.zeros(1, dtype=np.int32),\n        }\n\n        # Network parameters\n        for net_type_key in networks.keys():\n            for agent_net_key in networks[net_type_key].keys():\n                # Ensure obs and target networks are sonnet modules\n                server.store.parameters[f\"{net_type_key}-{agent_net_key}\"] = networks[\n                    net_type_key\n                ][agent_net_key].params\n\n        # Create the checkpointer\n        if self.config.checkpoint:\n            server.store.last_checkpoint_time = 0\n\n            # Only save variables that are not empty.\n            save_variables = {}\n            for key in server.store.parameters.keys():\n                var = server.store.parameters[key]\n                # Don't store empty tuple (e.g. empty observation_network) variables\n                if not (type(var) == tuple and len(var) == 0):\n                    save_variables[key] = var\n            server.store.system_checkpointer = savers.Checkpointer(\n                save_variables, self.config.checkpoint_subpath, time_delta_minutes=0\n            )\n\n    # Get\n    def on_parameter_server_get_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        names: Union[str, Sequence[str]] = server.store._param_names\n\n        if type(names) == str:\n            get_params = server.store.parameters[names]  # type: ignore\n        else:\n            get_params = {}\n            for var_key in names:\n                get_params[var_key] = server.store.parameters[var_key]\n        server.store.get_parameters = get_params\n\n    # Set\n    def on_parameter_server_set_parameters(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\"\"\"\n        params: Dict[str, Any] = server.store._set_params\n        names = params.keys()\n\n        for var_key in names:\n            assert var_key in server.store.parameters\n            if type(server.store.parameters[var_key]) == tuple:\n                raise NotImplementedError\n                # # Loop through tuple\n                # for var_i in range(len(server.store.parameters[var_key])):\n                #     server.store.parameters[var_key][var_i].assign(params[var_key][var_i])\n            else:\n                server.store.parameters[var_key] = params[var_key]\n\n    # Add\n    def on_parameter_server_add_to_parameters(\n        self, server: SystemParameterServer\n    ) -> None:\n        \"\"\"_summary_\"\"\"\n        params: Dict[str, Any] = server.store._add_to_params\n        names = params.keys()\n\n        for var_key in names:\n            assert var_key in server.store.parameters\n            server.store.parameters[var_key] += params[var_key]\n\n    # Save variables using checkpointer\n    def on_parameter_server_run_loop(self, server: SystemParameterServer) -> None:\n        \"\"\"_summary_\n\n        Args:\n            server : _description_\n        \"\"\"\n        if (\n            self.config.checkpoint\n            and server.store.last_checkpoint_time\n            + self.config.checkpoint_minute_interval * 60\n            + 1\n            < time.time()\n        ):\n            server.store.system_checkpointer.save()\n            server.store.last_checkpoint_time = time.time()\n            print(\"Updated variables checkpoint.\")\n", "path": "mava/components/jax/updating/parameter_server.py"}]}