We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[TEST] Jax MAPPO System Integration
### What do you want to test?
Jax MAPPO full integration test

### Outline of test structure
* Unit tests (if possible)
* Test component interactions
* Ensure not crashing during standard executor and trainer steps

### Definition of done
Passing checks, cover all basic component interactions, 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/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         lp_launch_type: Union[
47             str, lp.LaunchType
48         ] = lp.LaunchType.LOCAL_MULTI_PROCESSING,
49     ) -> None:
50         """Initialise the launcher.
51 
52         If multi-process, set up the launchpad program.
53         Otherwise, create a dictionary for the nodes in the system.
54 
55         Args:
56             multi_process : whether to use launchpad to run nodes on separate processes.
57             nodes_on_gpu : which nodes should be run on the GPU.
58             single_process_trainer_period : number of episodes between single process
59                 trainer steps.
60             single_process_evaluator_period : num episodes between single process
61                 evaluator steps.
62             single_process_max_episodes: maximum number of episodes to run
63                 before termination.
64             name : launchpad program name.
65             terminal : terminal for launchpad processes to be shown on.
66             lp_launch_type: launchpad launch type.
67         """
68         self._multi_process = multi_process
69         self._name = name
70         self._single_process_trainer_period = single_process_trainer_period
71         self._single_process_evaluator_period = single_process_evaluator_period
72         self._single_process_max_episodes = single_process_max_episodes
73         self._terminal = terminal
74         self._lp_launch_type = lp_launch_type
75         if multi_process:
76             self._program = lp.Program(name=name)
77             self._nodes_on_gpu = nodes_on_gpu
78         else:
79             self._nodes: List = []
80             self._node_dict: Dict = {
81                 "data_server": None,
82                 "parameter_server": None,
83                 "executor": None,
84                 "evaluator": None,
85                 "trainer": None,
86             }
87 
88     def add(
89         self,
90         node_fn: Any,
91         arguments: Any = [],
92         node_type: Union[lp.ReverbNode, lp.CourierNode] = NodeType.courier,
93         name: str = "Node",
94     ) -> Any:
95         """Add a node to the system.
96 
97         If multi-processing, add a node to the existing launchpad program,
98         grouped under the given name.
99         This means that when multi-processing,
100         you can have multiple nodes of the same name (e.g. executor).
101         If system is single-process, only one node per name is allowed in the system.
102 
103         Args:
104             node_fn : Function returning the system process that will run on the node.
105             arguments : Arguments used when initialising the system process.
106             node_type : Type of launchpad node to use.
107             name : Node name (e.g. executor).
108 
109         Raises:
110             ValueError: if single-process and node name is not supported.
111             ValueError: if single-process and trying to init a node more than once.
112 
113         Returns:
114             The system process or launchpad node.
115         """
116         # Create a list of arguments
117         if type(arguments) is not list:
118             arguments = [arguments]
119 
120         if self._multi_process:
121             with self._program.group(name):
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             local_resources = lp_utils.to_device(
169                 program_nodes=self._program.groups.keys(),
170                 nodes_on_gpu=self._nodes_on_gpu,
171             )
172 
173             lp.launch(
174                 self._program,
175                 launch_type=self._lp_launch_type,
176                 terminal=self._terminal,
177                 local_resources=local_resources,
178             )
179         else:
180             episode = 1
181             step = 1
182             executor_steps = 0
183 
184             data_server = self._node_dict["data_server"]
185             _ = self._node_dict["parameter_server"]
186             executor = self._node_dict["executor"]
187             evaluator = self._node_dict["evaluator"]
188             trainer = self._node_dict["trainer"]
189 
190             # getting the maximum queue size
191             queue_threshold = data_server.server_info()["trainer"].max_size
192 
193             while (
194                 self._single_process_max_episodes is None
195                 or episode <= self._single_process_max_episodes
196             ):
197                 # if the queue is too full we skip the executor to ensure that the
198                 # executor won't hang when trying to push experience
199                 if data_server.server_info()["trainer"].current_size < int(
200                     queue_threshold * 0.75
201                 ):
202                     executor_stats = executor.run_episode_and_log()
203                     executor_steps += executor_stats["episode_length"]
204 
205                     print(f"Episode {episode} completed.")
206                     episode += 1
207 
208                 # if the queue has less than sample_batch_size samples in it we skip
209                 # the trainer to ensure that the trainer won't hang
210                 if (
211                     data_server.server_info()["trainer"].current_size
212                     >= trainer.store.global_config.sample_batch_size
213                     and step % self._single_process_trainer_period == 0
214                 ):
215                     _ = trainer.step()  # logging done in trainer
216                     print("Performed trainer step.")
217                 if step % self._single_process_evaluator_period == 0:
218                     _ = evaluator.run_episode_and_log()
219                     print("Performed evaluator run.")
220 
221                 step += 1
222
```
--- 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
@@ -119,6 +119,7 @@
 
         if self._multi_process:
             with self._program.group(name):
+                node_fn = copy_node_fn(node_fn)
                 node = self._program.add_node(node_type(node_fn, *arguments))
             return node
         else:

{"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@@ -119,6 +119,7 @@\n \n         if self._multi_process:\n             with self._program.group(name):\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", "issue": "[TEST] Jax MAPPO System Integration\n### What do you want to test?\r\nJax MAPPO full integration test\r\n\r\n### Outline of test structure\r\n* Unit tests (if possible)\r\n* Test component interactions\r\n* Ensure not crashing during standard executor and trainer steps\r\n\r\n### Definition of done\r\nPassing checks, cover all basic component interactions, 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\"\"\"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        lp_launch_type: Union[\n            str, lp.LaunchType\n        ] = lp.LaunchType.LOCAL_MULTI_PROCESSING,\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            lp_launch_type: launchpad launch type.\n        \"\"\"\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        self._lp_launch_type = lp_launch_type\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                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            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=self._lp_launch_type,\n                terminal=self._terminal,\n                local_resources=local_resources,\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        lp_launch_type: Union[\n            str, lp.LaunchType\n        ] = lp.LaunchType.LOCAL_MULTI_PROCESSING,\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            lp_launch_type: launchpad launch type.\n        \"\"\"\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        self._lp_launch_type = lp_launch_type\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                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            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=self._lp_launch_type,\n                terminal=self._terminal,\n                local_resources=local_resources,\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"}]}