We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Single Process Example Fails
### Describe the bug
Error:
```
Traceback (most recent call last):
  File "/home/kaleab/Documents/Code/Mava/Mava/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py", line 110, in <module>
    app.run(main)
  File "/home/kaleab/anaconda3/envs/mava/lib/python3.9/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/home/kaleab/anaconda3/envs/mava/lib/python3.9/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
  File "/home/kaleab/Documents/Code/Mava/Mava/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py", line 106, in main
    system.launch()
  File "/home/kaleab/Documents/Code/Mava/Mava/mava/systems/jax/system.py", line 185, in launch
    self._builder.launch()
  File "/home/kaleab/Documents/Code/Mava/Mava/mava/systems/jax/builder.py", line 244, in launch
    self.on_building_launch()
  File "/home/kaleab/Documents/Code/Mava/Mava/mava/callbacks/builder_mixin.py", line 190, in on_building_launch
    callback.on_building_launch(self)
  File "/home/kaleab/Documents/Code/Mava/Mava/mava/components/jax/building/distributor.py", line 147, in on_building_launch
    builder.store.program.launch()
  File "/home/kaleab/Documents/Code/Mava/Mava/mava/systems/jax/launcher.py", line 197, in launch
    queue_threshold = data_server.server_info()["trainer"].max_size
KeyError: 'trainer'
[reverb/cc/platform/default/server.cc:84] Shutting down replay server
```

### To Reproduce
Steps to reproduce the behavior:
1. `python examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py`


### Expected behavior
Example to run. 

### Context (Environment)
 - Mava `dev`.

### Additional context
-

### Possible Solution
-


--- 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/systems/jax/launcher.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 """General launcher for systems"""
17 from typing import Any, Dict, List, Optional, Union
18 
19 import launchpad as lp
20 import reverb
21 
22 from mava.utils import lp_utils
23 from mava.utils.builder_utils import copy_node_fn
24 
25 
26 class NodeType:
27     """Specify launchpad node types that systems can use."""
28 
29     reverb = lp.ReverbNode
30     courier = lp.CourierNode
31 
32 
33 class Launcher:
34     """This mava launcher can be used to launch multi-node systems using either single \
35         or distributed computation."""
36 
37     def __init__(
38         self,
39         multi_process: bool,
40         nodes_on_gpu: List = [],
41         single_process_trainer_period: int = 1,
42         single_process_evaluator_period: int = 10,
43         single_process_max_episodes: Optional[int] = None,
44         name: str = "System",
45         terminal: str = "current_terminal",
46         is_test: Optional[bool] = False,
47     ) -> None:
48         """Initialise the launcher.
49 
50         If multi-process, set up the launchpad program.
51         Otherwise, create a dictionary for the nodes in the system.
52 
53         Args:
54             multi_process : whether to use launchpad to run nodes on separate processes.
55             nodes_on_gpu : which nodes should be run on the GPU.
56             single_process_trainer_period : number of episodes between single process
57                 trainer steps.
58             single_process_evaluator_period : num episodes between single process
59                 evaluator steps.
60             single_process_max_episodes: maximum number of episodes to run
61                 before termination.
62             name : launchpad program name.
63             terminal : terminal for launchpad processes to be shown on.
64             is_test : whether to set testing launchpad launch_type.
65         """
66         self._is_test = is_test
67         self._multi_process = multi_process
68         self._name = name
69         self._single_process_trainer_period = single_process_trainer_period
70         self._single_process_evaluator_period = single_process_evaluator_period
71         self._single_process_max_episodes = single_process_max_episodes
72         self._terminal = terminal
73         if multi_process:
74             self._program = lp.Program(name=name)
75             self._nodes_on_gpu = nodes_on_gpu
76         else:
77             self._nodes: List = []
78             self._node_dict: Dict = {
79                 "data_server": None,
80                 "parameter_server": None,
81                 "executor": None,
82                 "evaluator": None,
83                 "trainer": None,
84             }
85 
86     def add(
87         self,
88         node_fn: Any,
89         arguments: Any = [],
90         node_type: Union[lp.ReverbNode, lp.CourierNode] = NodeType.courier,
91         name: str = "Node",
92     ) -> Any:
93         """Add a node to the system.
94 
95         If multi-processing, add a node to the existing launchpad program,
96         grouped under the given name.
97         This means that when multi-processing,
98         you can have multiple nodes of the same name (e.g. executor).
99         If system is single-process, only one node per name is allowed in the system.
100 
101         Args:
102             node_fn : Function returning the system process that will run on the node.
103             arguments : Arguments used when initialising the system process.
104             node_type : Type of launchpad node to use.
105             name : Node name (e.g. executor).
106 
107         Raises:
108             ValueError: if single-process and node name is not supported.
109             ValueError: if single-process and trying to init a node more than once.
110 
111         Returns:
112             The system process or launchpad node.
113         """
114         # Create a list of arguments
115         if type(arguments) is not list:
116             arguments = [arguments]
117 
118         if self._multi_process:
119             with self._program.group(name):
120                 if self._is_test:
121                     node_fn = copy_node_fn(node_fn)
122                 node = self._program.add_node(node_type(node_fn, *arguments))
123             return node
124         else:
125             if name not in self._node_dict:
126                 raise ValueError(
127                     f"{name} is not a valid node name."
128                     + "Single process currently only supports "
129                     + "nodes named: {list(self._node_dict.keys())}"
130                 )
131             elif self._node_dict[name] is not None:
132                 raise ValueError(
133                     f"Node named {name} initialised more than once."
134                     + "Single process currently only supports one node per type."
135                 )
136 
137             node_fn = copy_node_fn(node_fn)
138             process = node_fn(*arguments)
139             if node_type == lp.ReverbNode:
140                 # Assigning server to self to keep it alive.
141                 self._replay_server = reverb.Server(process, port=None)
142                 process = reverb.Client(f"localhost:{self._replay_server.port}")
143             self._nodes.append(process)
144             self._node_dict[name] = process
145             return process
146 
147     def get_nodes(self) -> List[Any]:
148         """Get the nodes of a single-process system.
149 
150         Raises:
151             ValueError: if system is multi-process.
152 
153         Returns:
154             System nodes.
155         """
156         if self._multi_process:
157             raise ValueError("Get nodes only implemented for single process setups.")
158 
159         return self._nodes
160 
161     def launch(self) -> None:
162         """Launch the launchpad program or start the single-process system loop.
163 
164         Returns:
165             None.
166         """
167         if self._multi_process:
168             if self._is_test:
169                 launch_type = lp.LaunchType.TEST_MULTI_THREADING
170             else:
171                 launch_type = lp.LaunchType.LOCAL_MULTI_PROCESSING
172 
173             local_resources = lp_utils.to_device(
174                 program_nodes=self._program.groups.keys(),
175                 nodes_on_gpu=self._nodes_on_gpu,
176             )
177 
178             lp.launch(
179                 self._program,
180                 launch_type=launch_type,
181                 terminal=self._terminal,
182                 local_resources=local_resources,
183             )
184 
185         else:
186             episode = 1
187             step = 1
188             executor_steps = 0
189 
190             data_server = self._node_dict["data_server"]
191             _ = self._node_dict["parameter_server"]
192             executor = self._node_dict["executor"]
193             evaluator = self._node_dict["evaluator"]
194             trainer = self._node_dict["trainer"]
195 
196             # getting the maximum queue size
197             queue_threshold = data_server.server_info()["trainer"].max_size
198 
199             while (
200                 self._single_process_max_episodes is None
201                 or episode <= self._single_process_max_episodes
202             ):
203                 # if the queue is too full we skip the executor to ensure that the
204                 # executor won't hang when trying to push experience
205                 if data_server.server_info()["trainer"].current_size < int(
206                     queue_threshold * 0.75
207                 ):
208                     executor_stats = executor.run_episode_and_log()
209                     executor_steps += executor_stats["episode_length"]
210 
211                     print(f"Episode {episode} completed.")
212                     episode += 1
213 
214                 # if the queue has less than sample_batch_size samples in it we skip
215                 # the trainer to ensure that the trainer won't hang
216                 if (
217                     data_server.server_info()["trainer"].current_size
218                     >= trainer.store.global_config.sample_batch_size
219                     and step % self._single_process_trainer_period == 0
220                 ):
221                     _ = trainer.step()  # logging done in trainer
222                     print("Performed trainer step.")
223                 if step % self._single_process_evaluator_period == 0:
224                     _ = evaluator.run_episode_and_log()
225                     print("Performed evaluator run.")
226 
227                 step += 1
228
```
--- 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/systems/jax/launcher.py b/mava/systems/jax/launcher.py
--- a/mava/systems/jax/launcher.py
+++ b/mava/systems/jax/launcher.py
@@ -194,7 +194,7 @@
             trainer = self._node_dict["trainer"]
 
             # getting the maximum queue size
-            queue_threshold = data_server.server_info()["trainer"].max_size
+            queue_threshold = data_server.server_info()["trainer_0"].max_size
 
             while (
                 self._single_process_max_episodes is None
@@ -202,7 +202,7 @@
             ):
                 # if the queue is too full we skip the executor to ensure that the
                 # executor won't hang when trying to push experience
-                if data_server.server_info()["trainer"].current_size < int(
+                if data_server.server_info()["trainer_0"].current_size < int(
                     queue_threshold * 0.75
                 ):
                     executor_stats = executor.run_episode_and_log()
@@ -214,7 +214,7 @@
                 # if the queue has less than sample_batch_size samples in it we skip
                 # the trainer to ensure that the trainer won't hang
                 if (
-                    data_server.server_info()["trainer"].current_size
+                    data_server.server_info()["trainer_0"].current_size
                     >= trainer.store.global_config.sample_batch_size
                     and step % self._single_process_trainer_period == 0
                 ):

{"golden_diff": "diff --git a/mava/systems/jax/launcher.py b/mava/systems/jax/launcher.py\n--- a/mava/systems/jax/launcher.py\n+++ b/mava/systems/jax/launcher.py\n@@ -194,7 +194,7 @@\n             trainer = self._node_dict[\"trainer\"]\n \n             # getting the maximum queue size\n-            queue_threshold = data_server.server_info()[\"trainer\"].max_size\n+            queue_threshold = data_server.server_info()[\"trainer_0\"].max_size\n \n             while (\n                 self._single_process_max_episodes is None\n@@ -202,7 +202,7 @@\n             ):\n                 # if the queue is too full we skip the executor to ensure that the\n                 # executor won't hang when trying to push experience\n-                if data_server.server_info()[\"trainer\"].current_size < int(\n+                if data_server.server_info()[\"trainer_0\"].current_size < int(\n                     queue_threshold * 0.75\n                 ):\n                     executor_stats = executor.run_episode_and_log()\n@@ -214,7 +214,7 @@\n                 # if the queue has less than sample_batch_size samples in it we skip\n                 # the trainer to ensure that the trainer won't hang\n                 if (\n-                    data_server.server_info()[\"trainer\"].current_size\n+                    data_server.server_info()[\"trainer_0\"].current_size\n                     >= trainer.store.global_config.sample_batch_size\n                     and step % self._single_process_trainer_period == 0\n                 ):\n", "issue": "[BUG] Single Process Example Fails\n### Describe the bug\r\nError:\r\n```\r\nTraceback (most recent call last):\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py\", line 110, in <module>\r\n    app.run(main)\r\n  File \"/home/kaleab/anaconda3/envs/mava/lib/python3.9/site-packages/absl/app.py\", line 308, in run\r\n    _run_main(main, args)\r\n  File \"/home/kaleab/anaconda3/envs/mava/lib/python3.9/site-packages/absl/app.py\", line 254, in _run_main\r\n    sys.exit(main(argv))\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py\", line 106, in main\r\n    system.launch()\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/mava/systems/jax/system.py\", line 185, in launch\r\n    self._builder.launch()\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/mava/systems/jax/builder.py\", line 244, in launch\r\n    self.on_building_launch()\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/mava/callbacks/builder_mixin.py\", line 190, in on_building_launch\r\n    callback.on_building_launch(self)\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/mava/components/jax/building/distributor.py\", line 147, in on_building_launch\r\n    builder.store.program.launch()\r\n  File \"/home/kaleab/Documents/Code/Mava/Mava/mava/systems/jax/launcher.py\", line 197, in launch\r\n    queue_threshold = data_server.server_info()[\"trainer\"].max_size\r\nKeyError: 'trainer'\r\n[reverb/cc/platform/default/server.cc:84] Shutting down replay server\r\n```\r\n\r\n### To Reproduce\r\nSteps to reproduce the behavior:\r\n1. `python examples/jax/debugging/simple_spread/feedforward/decentralised/run_ippo_single_process.py`\r\n\r\n\r\n### Expected behavior\r\nExample to run. \r\n\r\n### Context (Environment)\r\n - Mava `dev`.\r\n\r\n### Additional context\r\n-\r\n\r\n### Possible Solution\r\n-\r\n\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\"\"\"General launcher for systems\"\"\"\nfrom typing import Any, Dict, List, Optional, Union\n\nimport launchpad as lp\nimport reverb\n\nfrom mava.utils import lp_utils\nfrom mava.utils.builder_utils import copy_node_fn\n\n\nclass NodeType:\n    \"\"\"Specify launchpad node types that systems can use.\"\"\"\n\n    reverb = lp.ReverbNode\n    courier = lp.CourierNode\n\n\nclass Launcher:\n    \"\"\"This mava launcher can be used to launch multi-node systems using either single \\\n        or distributed computation.\"\"\"\n\n    def __init__(\n        self,\n        multi_process: bool,\n        nodes_on_gpu: List = [],\n        single_process_trainer_period: int = 1,\n        single_process_evaluator_period: int = 10,\n        single_process_max_episodes: Optional[int] = None,\n        name: str = \"System\",\n        terminal: str = \"current_terminal\",\n        is_test: Optional[bool] = False,\n    ) -> None:\n        \"\"\"Initialise the launcher.\n\n        If multi-process, set up the launchpad program.\n        Otherwise, create a dictionary for the nodes in the system.\n\n        Args:\n            multi_process : whether to use launchpad to run nodes on separate processes.\n            nodes_on_gpu : which nodes should be run on the GPU.\n            single_process_trainer_period : number of episodes between single process\n                trainer steps.\n            single_process_evaluator_period : num episodes between single process\n                evaluator steps.\n            single_process_max_episodes: maximum number of episodes to run\n                before termination.\n            name : launchpad program name.\n            terminal : terminal for launchpad processes to be shown on.\n            is_test : whether to set testing launchpad launch_type.\n        \"\"\"\n        self._is_test = is_test\n        self._multi_process = multi_process\n        self._name = name\n        self._single_process_trainer_period = single_process_trainer_period\n        self._single_process_evaluator_period = single_process_evaluator_period\n        self._single_process_max_episodes = single_process_max_episodes\n        self._terminal = terminal\n        if multi_process:\n            self._program = lp.Program(name=name)\n            self._nodes_on_gpu = nodes_on_gpu\n        else:\n            self._nodes: List = []\n            self._node_dict: Dict = {\n                \"data_server\": None,\n                \"parameter_server\": None,\n                \"executor\": None,\n                \"evaluator\": None,\n                \"trainer\": None,\n            }\n\n    def add(\n        self,\n        node_fn: Any,\n        arguments: Any = [],\n        node_type: Union[lp.ReverbNode, lp.CourierNode] = NodeType.courier,\n        name: str = \"Node\",\n    ) -> Any:\n        \"\"\"Add a node to the system.\n\n        If multi-processing, add a node to the existing launchpad program,\n        grouped under the given name.\n        This means that when multi-processing,\n        you can have multiple nodes of the same name (e.g. executor).\n        If system is single-process, only one node per name is allowed in the system.\n\n        Args:\n            node_fn : Function returning the system process that will run on the node.\n            arguments : Arguments used when initialising the system process.\n            node_type : Type of launchpad node to use.\n            name : Node name (e.g. executor).\n\n        Raises:\n            ValueError: if single-process and node name is not supported.\n            ValueError: if single-process and trying to init a node more than once.\n\n        Returns:\n            The system process or launchpad node.\n        \"\"\"\n        # Create a list of arguments\n        if type(arguments) is not list:\n            arguments = [arguments]\n\n        if self._multi_process:\n            with self._program.group(name):\n                if self._is_test:\n                    node_fn = copy_node_fn(node_fn)\n                node = self._program.add_node(node_type(node_fn, *arguments))\n            return node\n        else:\n            if name not in self._node_dict:\n                raise ValueError(\n                    f\"{name} is not a valid node name.\"\n                    + \"Single process currently only supports \"\n                    + \"nodes named: {list(self._node_dict.keys())}\"\n                )\n            elif self._node_dict[name] is not None:\n                raise ValueError(\n                    f\"Node named {name} initialised more than once.\"\n                    + \"Single process currently only supports one node per type.\"\n                )\n\n            node_fn = copy_node_fn(node_fn)\n            process = node_fn(*arguments)\n            if node_type == lp.ReverbNode:\n                # Assigning server to self to keep it alive.\n                self._replay_server = reverb.Server(process, port=None)\n                process = reverb.Client(f\"localhost:{self._replay_server.port}\")\n            self._nodes.append(process)\n            self._node_dict[name] = process\n            return process\n\n    def get_nodes(self) -> List[Any]:\n        \"\"\"Get the nodes of a single-process system.\n\n        Raises:\n            ValueError: if system is multi-process.\n\n        Returns:\n            System nodes.\n        \"\"\"\n        if self._multi_process:\n            raise ValueError(\"Get nodes only implemented for single process setups.\")\n\n        return self._nodes\n\n    def launch(self) -> None:\n        \"\"\"Launch the launchpad program or start the single-process system loop.\n\n        Returns:\n            None.\n        \"\"\"\n        if self._multi_process:\n            if self._is_test:\n                launch_type = lp.LaunchType.TEST_MULTI_THREADING\n            else:\n                launch_type = lp.LaunchType.LOCAL_MULTI_PROCESSING\n\n            local_resources = lp_utils.to_device(\n                program_nodes=self._program.groups.keys(),\n                nodes_on_gpu=self._nodes_on_gpu,\n            )\n\n            lp.launch(\n                self._program,\n                launch_type=launch_type,\n                terminal=self._terminal,\n                local_resources=local_resources,\n            )\n\n        else:\n            episode = 1\n            step = 1\n            executor_steps = 0\n\n            data_server = self._node_dict[\"data_server\"]\n            _ = self._node_dict[\"parameter_server\"]\n            executor = self._node_dict[\"executor\"]\n            evaluator = self._node_dict[\"evaluator\"]\n            trainer = self._node_dict[\"trainer\"]\n\n            # getting the maximum queue size\n            queue_threshold = data_server.server_info()[\"trainer\"].max_size\n\n            while (\n                self._single_process_max_episodes is None\n                or episode <= self._single_process_max_episodes\n            ):\n                # if the queue is too full we skip the executor to ensure that the\n                # executor won't hang when trying to push experience\n                if data_server.server_info()[\"trainer\"].current_size < int(\n                    queue_threshold * 0.75\n                ):\n                    executor_stats = executor.run_episode_and_log()\n                    executor_steps += executor_stats[\"episode_length\"]\n\n                    print(f\"Episode {episode} completed.\")\n                    episode += 1\n\n                # if the queue has less than sample_batch_size samples in it we skip\n                # the trainer to ensure that the trainer won't hang\n                if (\n                    data_server.server_info()[\"trainer\"].current_size\n                    >= trainer.store.global_config.sample_batch_size\n                    and step % self._single_process_trainer_period == 0\n                ):\n                    _ = trainer.step()  # logging done in trainer\n                    print(\"Performed trainer step.\")\n                if step % self._single_process_evaluator_period == 0:\n                    _ = evaluator.run_episode_and_log()\n                    print(\"Performed evaluator run.\")\n\n                step += 1\n", "path": "mava/systems/jax/launcher.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\"\"\"General launcher for systems\"\"\"\nfrom typing import Any, Dict, List, Optional, Union\n\nimport launchpad as lp\nimport reverb\n\nfrom mava.utils import lp_utils\nfrom mava.utils.builder_utils import copy_node_fn\n\n\nclass NodeType:\n    \"\"\"Specify launchpad node types that systems can use.\"\"\"\n\n    reverb = lp.ReverbNode\n    courier = lp.CourierNode\n\n\nclass Launcher:\n    \"\"\"This mava launcher can be used to launch multi-node systems using either single \\\n        or distributed computation.\"\"\"\n\n    def __init__(\n        self,\n        multi_process: bool,\n        nodes_on_gpu: List = [],\n        single_process_trainer_period: int = 1,\n        single_process_evaluator_period: int = 10,\n        single_process_max_episodes: Optional[int] = None,\n        name: str = \"System\",\n        terminal: str = \"current_terminal\",\n        is_test: Optional[bool] = False,\n    ) -> None:\n        \"\"\"Initialise the launcher.\n\n        If multi-process, set up the launchpad program.\n        Otherwise, create a dictionary for the nodes in the system.\n\n        Args:\n            multi_process : whether to use launchpad to run nodes on separate processes.\n            nodes_on_gpu : which nodes should be run on the GPU.\n            single_process_trainer_period : number of episodes between single process\n                trainer steps.\n            single_process_evaluator_period : num episodes between single process\n                evaluator steps.\n            single_process_max_episodes: maximum number of episodes to run\n                before termination.\n            name : launchpad program name.\n            terminal : terminal for launchpad processes to be shown on.\n            is_test : whether to set testing launchpad launch_type.\n        \"\"\"\n        self._is_test = is_test\n        self._multi_process = multi_process\n        self._name = name\n        self._single_process_trainer_period = single_process_trainer_period\n        self._single_process_evaluator_period = single_process_evaluator_period\n        self._single_process_max_episodes = single_process_max_episodes\n        self._terminal = terminal\n        if multi_process:\n            self._program = lp.Program(name=name)\n            self._nodes_on_gpu = nodes_on_gpu\n        else:\n            self._nodes: List = []\n            self._node_dict: Dict = {\n                \"data_server\": None,\n                \"parameter_server\": None,\n                \"executor\": None,\n                \"evaluator\": None,\n                \"trainer\": None,\n            }\n\n    def add(\n        self,\n        node_fn: Any,\n        arguments: Any = [],\n        node_type: Union[lp.ReverbNode, lp.CourierNode] = NodeType.courier,\n        name: str = \"Node\",\n    ) -> Any:\n        \"\"\"Add a node to the system.\n\n        If multi-processing, add a node to the existing launchpad program,\n        grouped under the given name.\n        This means that when multi-processing,\n        you can have multiple nodes of the same name (e.g. executor).\n        If system is single-process, only one node per name is allowed in the system.\n\n        Args:\n            node_fn : Function returning the system process that will run on the node.\n            arguments : Arguments used when initialising the system process.\n            node_type : Type of launchpad node to use.\n            name : Node name (e.g. executor).\n\n        Raises:\n            ValueError: if single-process and node name is not supported.\n            ValueError: if single-process and trying to init a node more than once.\n\n        Returns:\n            The system process or launchpad node.\n        \"\"\"\n        # Create a list of arguments\n        if type(arguments) is not list:\n            arguments = [arguments]\n\n        if self._multi_process:\n            with self._program.group(name):\n                if self._is_test:\n                    node_fn = copy_node_fn(node_fn)\n                node = self._program.add_node(node_type(node_fn, *arguments))\n            return node\n        else:\n            if name not in self._node_dict:\n                raise ValueError(\n                    f\"{name} is not a valid node name.\"\n                    + \"Single process currently only supports \"\n                    + \"nodes named: {list(self._node_dict.keys())}\"\n                )\n            elif self._node_dict[name] is not None:\n                raise ValueError(\n                    f\"Node named {name} initialised more than once.\"\n                    + \"Single process currently only supports one node per type.\"\n                )\n\n            node_fn = copy_node_fn(node_fn)\n            process = node_fn(*arguments)\n            if node_type == lp.ReverbNode:\n                # Assigning server to self to keep it alive.\n                self._replay_server = reverb.Server(process, port=None)\n                process = reverb.Client(f\"localhost:{self._replay_server.port}\")\n            self._nodes.append(process)\n            self._node_dict[name] = process\n            return process\n\n    def get_nodes(self) -> List[Any]:\n        \"\"\"Get the nodes of a single-process system.\n\n        Raises:\n            ValueError: if system is multi-process.\n\n        Returns:\n            System nodes.\n        \"\"\"\n        if self._multi_process:\n            raise ValueError(\"Get nodes only implemented for single process setups.\")\n\n        return self._nodes\n\n    def launch(self) -> None:\n        \"\"\"Launch the launchpad program or start the single-process system loop.\n\n        Returns:\n            None.\n        \"\"\"\n        if self._multi_process:\n            if self._is_test:\n                launch_type = lp.LaunchType.TEST_MULTI_THREADING\n            else:\n                launch_type = lp.LaunchType.LOCAL_MULTI_PROCESSING\n\n            local_resources = lp_utils.to_device(\n                program_nodes=self._program.groups.keys(),\n                nodes_on_gpu=self._nodes_on_gpu,\n            )\n\n            lp.launch(\n                self._program,\n                launch_type=launch_type,\n                terminal=self._terminal,\n                local_resources=local_resources,\n            )\n\n        else:\n            episode = 1\n            step = 1\n            executor_steps = 0\n\n            data_server = self._node_dict[\"data_server\"]\n            _ = self._node_dict[\"parameter_server\"]\n            executor = self._node_dict[\"executor\"]\n            evaluator = self._node_dict[\"evaluator\"]\n            trainer = self._node_dict[\"trainer\"]\n\n            # getting the maximum queue size\n            queue_threshold = data_server.server_info()[\"trainer_0\"].max_size\n\n            while (\n                self._single_process_max_episodes is None\n                or episode <= self._single_process_max_episodes\n            ):\n                # if the queue is too full we skip the executor to ensure that the\n                # executor won't hang when trying to push experience\n                if data_server.server_info()[\"trainer_0\"].current_size < int(\n                    queue_threshold * 0.75\n                ):\n                    executor_stats = executor.run_episode_and_log()\n                    executor_steps += executor_stats[\"episode_length\"]\n\n                    print(f\"Episode {episode} completed.\")\n                    episode += 1\n\n                # if the queue has less than sample_batch_size samples in it we skip\n                # the trainer to ensure that the trainer won't hang\n                if (\n                    data_server.server_info()[\"trainer_0\"].current_size\n                    >= trainer.store.global_config.sample_batch_size\n                    and step % self._single_process_trainer_period == 0\n                ):\n                    _ = trainer.step()  # logging done in trainer\n                    print(\"Performed trainer step.\")\n                if step % self._single_process_evaluator_period == 0:\n                    _ = evaluator.run_episode_and_log()\n                    print(\"Performed evaluator run.\")\n\n                step += 1\n", "path": "mava/systems/jax/launcher.py"}]}