Repository: instadeepai/Mava

Issue:
[FEATURE] custom Mava config class
### Feature
A custom class to handle config for Mava systems.

### Proposal
The class should be easy to use for setting system configs and be able to grow/shrink dynamically based on the component configuration of the system. The proposal is to achieve this using `SimpleNamespaces` containing a `dataclass` for each component. The class will also support a property method to easily retrieve the config for the system and a method for specific component configs by name. 

### Testing
All methods should be tested including specific get 
As input, the config will expect a dataclass with side effects of updating the current config.
The "get" method for component-specific config should take in the component name and return that component's config.


### Definition of done
The class is implemented and supports the above functionality.

### Mandatory checklist before making a PR
* [ ] The success criteria laid down in “Definition of done” are met.
* [ ] Code is documented - docstrings for methods and classes, static types for arguments.
* [ ] Code is tested - unit, integration and/or functional tests are added.
* [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.
* [ ] All functional tests are green.


Discussion:

Resolving patch:
diff --git a/mava/systems/jax/__init__.py b/mava/systems/jax/__init__.py
new file mode 100644
--- /dev/null
+++ b/mava/systems/jax/__init__.py
@@ -0,0 +1,18 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Imports for Jax-based Mava systems"""
+
+from mava.systems.jax.config import Config
diff --git a/mava/systems/jax/config.py b/mava/systems/jax/config.py
new file mode 100644
--- /dev/null
+++ b/mava/systems/jax/config.py
@@ -0,0 +1,153 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Config class for Mava systems"""
+
+from dataclasses import is_dataclass
+from types import SimpleNamespace
+from typing import Any, Dict, List
+
+from mava.utils.config_utils import flatten_dict
+
+
+class Config:
+    """Config handler for Jax-based Mava systems."""
+
+    def __init__(self) -> None:
+        """Initialise config"""
+        self._config: Dict = {}
+        self._current_params: List = []
+        self._built = False
+
+    def add(self, **kwargs: Any) -> None:
+        """Add a component config dataclass.
+
+        Raises:
+            Exception: if a config for an identically named component already exists
+            Exception: if a config shares a parameter name with another config
+            Exception: if a config is not a dataclass object
+        """
+        if not self._built:
+            for name, dataclass in kwargs.items():
+                if is_dataclass(dataclass):
+                    if name in list(self._config.keys()):
+                        raise Exception(
+                            "The given component config is already part of the current \
+                            system. Perhaps try updating the component instead using \
+                            .update() in the system builder."
+                        )
+                    else:
+                        new_param_names = list(dataclass.__dict__.keys())
+                        if set(self._current_params) & set(new_param_names):
+                            raise Exception(
+                                "Component configs share a common parameter name. \
+                                This is not allowed, please ensure config \
+                                parameter names are unique."
+                            )
+                        else:
+                            self._current_params.extend(new_param_names)
+                        self._config[name] = dataclass
+                else:
+                    raise Exception("Component configs must be a dataclass.")
+        else:
+            raise Exception(
+                "Component configs cannot be added to an already built config."
+            )
+
+    def update(self, **kwargs: Any) -> None:
+        """Update a component config dataclass.
+
+        Raises:
+            Exception: if a config shares a parameter name with another config
+            Exception: if a config is not already part of the system
+            Exception: if a config is not a dataclass object
+        """
+        if not self._built:
+            for name, dataclass in kwargs.items():
+                if is_dataclass(dataclass):
+                    if name in list(self._config.keys()):
+                        new_param_names = list(dataclass.__dict__.keys())
+                        if set(self._current_params) & set(new_param_names):
+                            raise Exception(
+                                "Component configs share a common parameter name. \
+                                This is not allowed, please ensure config \
+                                parameter names are unique."
+                            )
+                        else:
+                            self._current_params.extend(new_param_names)
+                        self._config[name] = dataclass
+                    else:
+                        raise Exception(
+                            "The given component config is not part of the current \
+                            system. Perhaps try adding the component using .add() \
+                            in the system builder."
+                        )
+                else:
+                    raise Exception("Component configs must be a dataclass.")
+        else:
+            raise Exception(
+                "Component configs cannot be updated if config has already been built."
+            )
+
+    def build(self) -> None:
+        """Build the config file, i.e. unwrap dataclass nested dictionaries"""
+        if not self._built:
+            config_unwrapped: Dict = {}
+            for param in self._config.values():
+                config_unwrapped.update(flatten_dict(param.__dict__))
+            self._built_config = config_unwrapped
+            self._built = True
+        else:
+            raise Exception("Config has already been built, this can only happen once.")
+
+    def set_parameters(self, **kwargs: Any) -> None:
+        """Set a specific hyperparameter of a built config.
+
+        Raises:
+            Exception: if a set is attempted on a config not yet built.
+            Exception: if a set is attempted for a hyperparameter that is not part \
+                of the built config.
+        """
+
+        if not self._built:
+            raise Exception(
+                "Config must first be built using .build() before hyperparameters \
+                can be set to different values using .set()."
+            )
+        else:
+            for name, param_value in kwargs.items():
+                if name in list(self._built_config.keys()):
+                    self._built_config[name] = param_value
+                else:
+                    raise Exception(
+                        "The given parameter is not part of the current system. \
+                        This should have been added first via a component .add() \
+                        during system building."
+                    )
+
+    def get(self) -> SimpleNamespace:
+        """Get built config for feeding to a Mava system.
+
+        Raises:
+            Exception: if trying to get without having first built the config
+        Returns:
+            built config
+        """
+        if self._built:
+            return SimpleNamespace(**self._built_config)
+        else:
+            raise Exception(
+                "The config must first be built using .build() before calling .get()."
+            )
diff --git a/mava/utils/config_utils.py b/mava/utils/config_utils.py
new file mode 100644
--- /dev/null
+++ b/mava/utils/config_utils.py
@@ -0,0 +1,41 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utility functions for Mava config"""
+
+from collections.abc import MutableMapping
+from typing import List
+
+
+def flatten_dict(
+    d: MutableMapping, parent_key: str = "", sep: str = "."
+) -> MutableMapping:
+    """Flattens nested dictionary structure.
+
+    Args:
+        d : dictionary to be flattened
+        parent_key : optional parent key identifier
+        sep : optional separator between parent and leaf keys
+    Returns:
+        Flattened dictionary
+    """
+    items: List = []
+    for k, v in d.items():
+        new_key = parent_key + sep + k if parent_key else k
+        if isinstance(v, MutableMapping):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)


Test patch:
diff --git a/mava/systems/jax/config_test.py b/mava/systems/jax/config_test.py
new file mode 100644
--- /dev/null
+++ b/mava/systems/jax/config_test.py
@@ -0,0 +1,250 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for config class for Jax-based Mava systems"""
+
+from dataclasses import dataclass
+
+import pytest
+
+from mava.systems.jax import Config
+
+
+@dataclass
+class ComponentConfig:
+    name: str
+    setting: int
+
+
+@dataclass
+class HyperparameterConfig:
+    param_0: float
+    param_1: float
+
+
+@dataclass
+class SameParameterNameConfig:
+    param_0: int
+    param_2: str
+
+
+@pytest.fixture
+def dummy_component_config() -> ComponentConfig:
+    """Dummy config dataclass for a component.
+
+    Returns:
+        config dataclass
+    """
+    return ComponentConfig(name="component", setting=5)
+
+
+@pytest.fixture
+def dummy_hyperparameter_config() -> HyperparameterConfig:
+    """Dummy config dataclass for component hyperparameters.
+
+    Returns:
+        config dataclass
+    """
+    return HyperparameterConfig(param_0=2.7, param_1=3.8)
+
+
+@pytest.fixture
+def config() -> Config:
+    """Config instance.
+
+    Returns:
+        instantiation of a Mava Config class
+    """
+    return Config()
+
+
+def test_add_single_config(config: Config, dummy_component_config: type) -> None:
+    """Test adding a single config.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+    """
+    config.add(component=dummy_component_config)
+    config.build()
+    conf = config.get()
+
+    assert conf.name == "component"
+    assert conf.setting == 5
+
+
+def test_add_multiple_configs(
+    config: Config, dummy_component_config: type, dummy_hyperparameter_config: type
+) -> None:
+    """Test adding multiple configs at the same time.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+        dummy_hyperparameter_config : component config dataclass of hyperparameters
+    """
+    config.add(
+        component=dummy_component_config, hyperparameter=dummy_hyperparameter_config
+    )
+    config.build()
+    conf = config.get()
+
+    assert conf.name == "component"
+    assert conf.setting == 5
+    assert conf.param_0 == 2.7
+    assert conf.param_1 == 3.8
+
+
+def test_add_configs_twice(
+    config: Config, dummy_component_config: type, dummy_hyperparameter_config: type
+) -> None:
+    """Test add two configs, one after the other.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+        dummy_hyperparameter_config : component config dataclass of hyperparameters
+    """
+    config.add(component=dummy_component_config)
+    config.add(hyperparameter=dummy_hyperparameter_config)
+    config.build()
+    conf = config.get()
+
+    assert conf.name == "component"
+    assert conf.setting == 5
+    assert conf.param_0 == 2.7
+    assert conf.param_1 == 3.8
+
+
+def test_update_config(
+    config: Config, dummy_component_config: type, dummy_hyperparameter_config: type
+) -> None:
+    """Test add two configs, one after the other.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+        dummy_hyperparameter_config : component config dataclass of hyperparameters
+    """
+    config.add(component=dummy_component_config)
+    config.update(component=dummy_hyperparameter_config)
+    config.build()
+    conf = config.get()
+
+    assert conf.param_0 == 2.7
+    assert conf.param_1 == 3.8
+    assert not hasattr(config, "name")
+    assert not hasattr(config, "setting")
+
+
+def test_set_existing_parameter_on_the_fly(
+    config: Config, dummy_component_config: type
+) -> None:
+    """Test updating a hyperparameter on the fly after the config has been built.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+    """
+
+    # add component dataclasses and build config
+    config.add(component=dummy_component_config)
+    config.build()
+
+    # set config parameters on the fly
+    config.set_parameters(name="new_component_name")
+    conf = config.get()
+
+    assert conf.name == "new_component_name"
+    assert conf.setting == 5
+
+
+def test_set_before_build_exception(
+    config: Config, dummy_component_config: type
+) -> None:
+    """Test that exception is thrown if it is attempted to set a hyperparameter \
+        before the config has been built.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+    """
+
+    with pytest.raises(Exception):
+        # add component dataclasses and build config
+        config.add(component=dummy_component_config)
+
+        # Try setting parameters without having built first
+        config.set_parameters(name="new_component_name")
+
+
+def test_get_before_build_exception(
+    config: Config, dummy_component_config: type
+) -> None:
+    """Test that exception is thrown if it is attempted to call .get() \
+        before the config has been built.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+    """
+
+    with pytest.raises(Exception):
+        # add component dataclasses and build config
+        config.add(component=dummy_component_config)
+
+        # Try getting without having built first
+        config.get()
+
+
+def test_parameter_setting_that_does_not_exist_exception(
+    config: Config, dummy_component_config: type
+) -> None:
+    """Test that exception is thrown if it is attempted to set a hyperparameter \
+        that does not exist.
+
+    Args:
+        config : Mava config
+        dummy_component_config : component config dataclass
+    """
+
+    with pytest.raises(Exception):
+        # add component dataclasses and build config
+        config.add(component=dummy_component_config)
+        config.build()
+
+        # Try setting a parameter that does not exist
+        config.set_parameters(unknown_param="new_value")
+
+
+def test_accidental_parameter_override_exception(
+    config: Config, dummy_hyperparameter_config: type
+) -> None:
+    """Test that exception is thrown when two component config dataclasses share the \
+        same name for a specific hyperparameter.
+
+    Args:
+        config : Mava config
+        dummy_hyperparameter_config : component config dataclass of hyperparameters
+    """
+
+    with pytest.raises(Exception):
+        # add component dataclasses and build config
+        config.add(hyperparameter=dummy_hyperparameter_config)
+
+        # add new component dataclass with a parameter of the same name
+        # as an already existing component parameter name
+        other_hyperparamter_config = SameParameterNameConfig(param_0=2, param_2="param")
+        config.add(other_hyperparameter=other_hyperparamter_config)
diff --git a/tests/utils/config_utils_test.py b/tests/utils/config_utils_test.py
new file mode 100644
--- /dev/null
+++ b/tests/utils/config_utils_test.py
@@ -0,0 +1,59 @@
+# python3
+# Copyright 2021 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import MutableMapping
+
+import pytest
+
+from mava.utils.config_utils import flatten_dict
+
+
+@pytest.fixture
+def test_mock_nested_dict() -> MutableMapping:
+    """Mock nested dictionary."""
+    return {"a": 1, "c": {"a": 2, "b": {"x": 5, "y": 10}}, "d": [1, 2, 3]}
+
+
+def test_flatten_dict(test_mock_nested_dict: MutableMapping) -> None:
+    """Tests flatten_dict."""
+    assert flatten_dict(d=test_mock_nested_dict) == {
+        "a": 1,
+        "c.a": 2,
+        "c.b.x": 5,
+        "c.b.y": 10,
+        "d": [1, 2, 3],
+    }
+
+
+def test_flatten_dict_different_sep(test_mock_nested_dict: MutableMapping) -> None:
+    """Tests flatten dict with specified sep token."""
+    assert flatten_dict(d=test_mock_nested_dict, sep="_") == {
+        "a": 1,
+        "c_a": 2,
+        "c_b_x": 5,
+        "c_b_y": 10,
+        "d": [1, 2, 3],
+    }
+
+
+def test_flatten_dict_with_parent_key(test_mock_nested_dict: MutableMapping) -> None:
+    """Tests flatten dict with specified parent key."""
+    assert flatten_dict(d=test_mock_nested_dict, parent_key="test_parent") == {
+        "test_parent.a": 1,
+        "test_parent.c.a": 2,
+        "test_parent.c.b.x": 5,
+        "test_parent.c.b.y": 10,
+        "test_parent.d": [1, 2, 3],
+    }