We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG][Jax] Default executor flattens observations
### Describe the bug
The default executor flattens observations which is undesirable for environments which can benefit from spatial knowledge.

### To Reproduce
`observation = executor.store.observation.observation.reshape((1, -1))` <- inside `action_selection.py`

### Expected behavior
It was expected that observations would be fed to neural networks unedited.

### Possible Solution
`observation = utils.add_batch_dim(executor.store.observation.observation)`


--- 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/executing/action_selection.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 """Execution components for system builders"""
17 
18 from dataclasses import dataclass
19 
20 import jax
21 
22 from mava.components.jax import Component
23 from mava.core_jax import SystemExecutor
24 
25 
26 @dataclass
27 class ExecutorSelectActionProcessConfig:
28     pass
29 
30 
31 class FeedforwardExecutorSelectAction(Component):
32     def __init__(
33         self,
34         config: ExecutorSelectActionProcessConfig = ExecutorSelectActionProcessConfig(),
35     ):
36         """_summary_
37 
38         Args:
39             config : _description_.
40         """
41         self.config = config
42 
43     # Select actions
44     def on_execution_select_actions(self, executor: SystemExecutor) -> None:
45         """Summary"""
46         executor.store.actions_info = {}
47         executor.store.policies_info = {}
48         for agent, observation in executor.store.observations.items():
49             action_info, policy_info = executor.select_action(agent, observation)
50             executor.store.actions_info[agent] = action_info
51             executor.store.policies_info[agent] = policy_info
52 
53     # Select action
54     def on_execution_select_action_compute(self, executor: SystemExecutor) -> None:
55         """Summary"""
56 
57         agent = executor.store.agent
58         network = executor.store.networks["networks"][
59             executor.store.agent_net_keys[agent]
60         ]
61 
62         observation = executor.store.observation.observation.reshape((1, -1))
63         rng_key, executor.store.key = jax.random.split(executor.store.key)
64 
65         # TODO (dries): We are currently using jit in the networks per agent.
66         # We can also try jit over all the agents in a for loop. This would
67         # allow the jit function to save us even more time.
68         executor.store.action_info, executor.store.policy_info = network.get_action(
69             observation, rng_key
70         )
71 
72     @staticmethod
73     def name() -> str:
74         """_summary_"""
75         return "action_selector"
76
```
--- 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/executing/action_selection.py b/mava/components/jax/executing/action_selection.py
--- a/mava/components/jax/executing/action_selection.py
+++ b/mava/components/jax/executing/action_selection.py
@@ -18,6 +18,7 @@
 from dataclasses import dataclass
 
 import jax
+from acme.jax import utils
 
 from mava.components.jax import Component
 from mava.core_jax import SystemExecutor
@@ -59,7 +60,7 @@
             executor.store.agent_net_keys[agent]
         ]
 
-        observation = executor.store.observation.observation.reshape((1, -1))
+        observation = utils.add_batch_dim(executor.store.observation.observation)
         rng_key, executor.store.key = jax.random.split(executor.store.key)
 
         # TODO (dries): We are currently using jit in the networks per agent.

{"golden_diff": "diff --git a/mava/components/jax/executing/action_selection.py b/mava/components/jax/executing/action_selection.py\n--- a/mava/components/jax/executing/action_selection.py\n+++ b/mava/components/jax/executing/action_selection.py\n@@ -18,6 +18,7 @@\n from dataclasses import dataclass\n \n import jax\n+from acme.jax import utils\n \n from mava.components.jax import Component\n from mava.core_jax import SystemExecutor\n@@ -59,7 +60,7 @@\n             executor.store.agent_net_keys[agent]\n         ]\n \n-        observation = executor.store.observation.observation.reshape((1, -1))\n+        observation = utils.add_batch_dim(executor.store.observation.observation)\n         rng_key, executor.store.key = jax.random.split(executor.store.key)\n \n         # TODO (dries): We are currently using jit in the networks per agent.\n", "issue": "[BUG][Jax] Default executor flattens observations\n### Describe the bug\r\nThe default executor flattens observations which is undesirable for environments which can benefit from spatial knowledge.\r\n\r\n### To Reproduce\r\n`observation = executor.store.observation.observation.reshape((1, -1))` <- inside `action_selection.py`\r\n\r\n### Expected behavior\r\nIt was expected that observations would be fed to neural networks unedited.\r\n\r\n### Possible Solution\r\n`observation = utils.add_batch_dim(executor.store.observation.observation)`\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\"\"\"Execution components for system builders\"\"\"\n\nfrom dataclasses import dataclass\n\nimport jax\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemExecutor\n\n\n@dataclass\nclass ExecutorSelectActionProcessConfig:\n    pass\n\n\nclass FeedforwardExecutorSelectAction(Component):\n    def __init__(\n        self,\n        config: ExecutorSelectActionProcessConfig = ExecutorSelectActionProcessConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    # Select actions\n    def on_execution_select_actions(self, executor: SystemExecutor) -> None:\n        \"\"\"Summary\"\"\"\n        executor.store.actions_info = {}\n        executor.store.policies_info = {}\n        for agent, observation in executor.store.observations.items():\n            action_info, policy_info = executor.select_action(agent, observation)\n            executor.store.actions_info[agent] = action_info\n            executor.store.policies_info[agent] = policy_info\n\n    # Select action\n    def on_execution_select_action_compute(self, executor: SystemExecutor) -> None:\n        \"\"\"Summary\"\"\"\n\n        agent = executor.store.agent\n        network = executor.store.networks[\"networks\"][\n            executor.store.agent_net_keys[agent]\n        ]\n\n        observation = executor.store.observation.observation.reshape((1, -1))\n        rng_key, executor.store.key = jax.random.split(executor.store.key)\n\n        # TODO (dries): We are currently using jit in the networks per agent.\n        # We can also try jit over all the agents in a for loop. This would\n        # allow the jit function to save us even more time.\n        executor.store.action_info, executor.store.policy_info = network.get_action(\n            observation, rng_key\n        )\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\"\"\"\n        return \"action_selector\"\n", "path": "mava/components/jax/executing/action_selection.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\"\"\"Execution components for system builders\"\"\"\n\nfrom dataclasses import dataclass\n\nimport jax\nfrom acme.jax import utils\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemExecutor\n\n\n@dataclass\nclass ExecutorSelectActionProcessConfig:\n    pass\n\n\nclass FeedforwardExecutorSelectAction(Component):\n    def __init__(\n        self,\n        config: ExecutorSelectActionProcessConfig = ExecutorSelectActionProcessConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    # Select actions\n    def on_execution_select_actions(self, executor: SystemExecutor) -> None:\n        \"\"\"Summary\"\"\"\n        executor.store.actions_info = {}\n        executor.store.policies_info = {}\n        for agent, observation in executor.store.observations.items():\n            action_info, policy_info = executor.select_action(agent, observation)\n            executor.store.actions_info[agent] = action_info\n            executor.store.policies_info[agent] = policy_info\n\n    # Select action\n    def on_execution_select_action_compute(self, executor: SystemExecutor) -> None:\n        \"\"\"Summary\"\"\"\n\n        agent = executor.store.agent\n        network = executor.store.networks[\"networks\"][\n            executor.store.agent_net_keys[agent]\n        ]\n\n        observation = utils.add_batch_dim(executor.store.observation.observation)\n        rng_key, executor.store.key = jax.random.split(executor.store.key)\n\n        # TODO (dries): We are currently using jit in the networks per agent.\n        # We can also try jit over all the agents in a for loop. This would\n        # allow the jit function to save us even more time.\n        executor.store.action_info, executor.store.policy_info = network.get_action(\n            observation, rng_key\n        )\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\"\"\"\n        return \"action_selector\"\n", "path": "mava/components/jax/executing/action_selection.py"}]}