Repository: instadeepai/Mava

Issue:
[FEATURE] Abstract builder class for Jax-based systems
### Feature
Abstract class for system building for Jax-based Mava systems.

### Proposal
The builder should take care of building essentially elements of a MARL system that run in different processes. These include the data server, variable server, executor (and evaluator) and trainer. 

### Testing
Tests will only consider the proper inheritance of the abstract builder class.

### Definition of done
All abstract methods are defined and have input and return types specified. 

### Mandatory checklist before making a PR
* [x] The success criteria laid down in “Definition of done” are met.
* [x] Code is documented - docstrings for methods and classes, static types for arguments.
* [x] Code is tested - unit, integration and/or functional tests are added.
* [x] Documentation is updated - README, CONTRIBUTING, or other documentation.
* [x] All functional tests are green.



Discussion:

Resolving patch:
diff --git a/mava/core_jax.py b/mava/core_jax.py
--- a/mava/core_jax.py
+++ b/mava/core_jax.py
@@ -69,3 +69,59 @@ def launch(
                 are primarily for debugging
             name : name of the system
         """
+
+
+class SystemBuilder(abc.ABC):
+    """Abstract system builder."""
+
+    @abc.abstractmethod
+    def data_server(self) -> List[Any]:
+        """Data server to store and serve transition data from and to system.
+
+        Returns:
+            System data server
+        """
+
+    @abc.abstractmethod
+    def parameter_server(self) -> Any:
+        """Parameter server to store and serve system network parameters.
+
+        Returns:
+            System parameter server
+        """
+
+    @abc.abstractmethod
+    def executor(
+        self, executor_id: str, data_server_client: Any, parameter_server_client: Any
+    ) -> Any:
+        """Executor, a collection of agents in an environment to gather experience.
+
+        Args:
+            executor_id : id to identify the executor process for logging purposes
+            data_server_client : data server client for pushing transition data
+            parameter_server_client : parameter server client for pulling parameters
+        Returns:
+            System executor
+        """
+
+    @abc.abstractmethod
+    def trainer(
+        self, trainer_id: str, data_server_client: Any, parameter_server_client: Any
+    ) -> Any:
+        """Trainer, a system process for updating agent specific network parameters.
+
+        Args:
+            trainer_id : id to identify the trainer process for logging purposes
+            data_server_client : data server client for pulling transition data
+            parameter_server_client : parameter server client for pushing parameters
+        Returns:
+            System trainer
+        """
+
+    @abc.abstractmethod
+    def build(self) -> None:
+        """Construct program nodes."""
+
+    @abc.abstractmethod
+    def launch(self) -> None:
+        """Run the graph program."""


Test patch:
diff --git a/mava/core_jax_test.py b/mava/core_jax_test.py
--- a/mava/core_jax_test.py
+++ b/mava/core_jax_test.py
@@ -16,14 +16,14 @@
 
 """Tests for core Mava interfaces for Jax systems."""
 
-from typing import Any
+from typing import Any, List
 
 import pytest
 
-from mava.core_jax import BaseSystem
+from mava.core_jax import BaseSystem, SystemBuilder
 
 
-def test_exception_for_incomplete_child_class() -> None:
+def test_exception_for_incomplete_child_system_class() -> None:
     """Test if error is thrown for missing abstract class overwrites."""
     with pytest.raises(TypeError):
 
@@ -41,3 +41,22 @@ def configure(self, **kwargs: Any) -> None:
                 pass
 
         TestIncompleteDummySystem()  # type: ignore
+
+
+def test_exception_for_incomplete_child_builder_class() -> None:
+    """Test if error is thrown for missing abstract class overwrites."""
+    with pytest.raises(TypeError):
+
+        class TestIncompleteDummySystemBuilder(SystemBuilder):
+            def data_server(self) -> List[Any]:
+                pass
+
+            def executor(
+                self,
+                executor_id: str,
+                data_server_client: Any,
+                parameter_server_client: Any,
+            ) -> Any:
+                pass
+
+        TestIncompleteDummySystemBuilder()  # type: ignore