Repository: instadeepai/Mava

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
-



Discussion:

Resolving patch:
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 @@ def launch(self) -> None:
             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 @@ def launch(self) -> None:
             ):
                 # 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 @@ def launch(self) -> None:
                 # 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
                 ):


Test patch: