Repository: instadeepai/Mava

Issue:
[MAINTAIN] Jax MAPPO documentation
### Please describe what needs to be maintained?
Jax MAPPO needs documentation: everything within `mava/systems/jax/mappo`

### Describe the outcome you'd like
All classes / methods / functions should be clearly documented with explanations and descriptions of all args / returns



Discussion:

Resolving patch:
diff --git a/mava/systems/jax/mappo/components.py b/mava/systems/jax/mappo/components.py
--- a/mava/systems/jax/mappo/components.py
+++ b/mava/systems/jax/mappo/components.py
@@ -33,19 +33,19 @@ class ExtrasLogProbSpecConfig:
 class ExtrasSpec(Component):
     @abc.abstractmethod
     def __init__(self, config: Any) -> None:
-        """_summary_
+        """Initialise extra specs
 
         Args:
-            config : _description_.
+            config : ExtrasSpecConfig
         """
         self.config = config
 
     @staticmethod
     def name() -> str:
-        """_summary_
+        """Returns name of ExtrasSpec class
 
         Returns:
-            _description_
+            "extras_spec": name of ExtrasSpec class
         """
         return "extras_spec"
 
@@ -55,15 +55,23 @@ def __init__(
         self,
         config: ExtrasLogProbSpecConfig = ExtrasLogProbSpecConfig(),
     ):
-        """_summary_
+        """Class that adds log probs to the extras spec
 
         Args:
-            config : _description_.
+            config : ExtrasLogProbSpecConfig
         """
         self.config = config
 
     def on_building_init_end(self, builder: SystemBuilder) -> None:
-        """[summary]"""
+        """Create extra specs after builder has been initialised
+
+        Args:
+            builder: SystemBuilder
+
+        Returns:
+            None.
+
+        """
         agent_specs = builder.store.ma_environment_spec.get_agent_environment_specs()
         builder.store.extras_spec = {"policy_info": {}}
 
diff --git a/mava/systems/jax/mappo/networks.py b/mava/systems/jax/mappo/networks.py
--- a/mava/systems/jax/mappo/networks.py
+++ b/mava/systems/jax/mappo/networks.py
@@ -39,7 +39,7 @@
 
 @dataclasses.dataclass
 class PPONetworks:
-    """TODO: Add description here."""
+    """Class to implement the networks for the PPO algorithm"""
 
     def __init__(
         self,
@@ -49,7 +49,15 @@ def __init__(
         entropy: Optional[EntropyFn] = None,
         sample: Optional[networks_lib.SampleFn] = None,
     ) -> None:
-        """TODO: Add description here."""
+        """Initialises the PPO network Class.
+
+        Args:
+            network: neural network representing the agent policy function.
+            params: values parameterising the network.
+            log_prob: function used to calculate the log prob of an agent's action.
+            entropy: function used to calculate the entropy of the agent policy.
+            sample: function used to select an action from the policy.
+        """
         self.network = network
         self.params = params
         self.log_prob = log_prob
@@ -63,7 +71,18 @@ def forward_fn(
             key: networks_lib.PRNGKey,
             mask: chex.Array = None,
         ) -> Tuple[jnp.ndarray, jnp.ndarray]:
-            """TODO: Add description here."""
+            """Forward function for PPO network
+
+            Args:
+                params: values parameterising the network.
+                observations: agent observations
+                key: pseudo-random value used to initialise distributions
+                mask: action mask which removes illegal actions
+
+            Returns:
+                actions: agent action
+                log_prob: log prob of the chosen action
+            """
             # The parameters of the network might change. So it has to
             # be fed into the jitted function.
             distribution, _ = self.network.apply(params, observations)
@@ -83,14 +102,33 @@ def get_action(
         key: networks_lib.PRNGKey,
         mask: chex.Array = None,
     ) -> Tuple[np.ndarray, Dict]:
-        """TODO: Add description here."""
+        """Gets an action from the network from given observation
+
+        Args:
+           observations: agent observations
+           key: pseudo-random value used to initialise distributions
+           mask: action mask which removes illegal actions
+
+        Returns:
+            actions: agent action
+            log_prob: log prob of the chosen action
+
+        """
         actions, log_prob = self.forward_fn(self.params, observations, key, mask)
         actions = np.array(actions, dtype=np.int64)
         log_prob = np.squeeze(np.array(log_prob, dtype=np.float32))
         return actions, {"log_prob": log_prob}
 
     def get_value(self, observations: networks_lib.Observation) -> jnp.ndarray:
-        """TODO: Add description here."""
+        """Gets value of observation
+
+        Args:
+            observations: agent observations
+
+        Returns:
+            value: estimated value of observation
+
+        """
         _, value = self.network.apply(self.params, observations)
         return value
 
@@ -98,7 +136,15 @@ def get_value(self, observations: networks_lib.Observation) -> jnp.ndarray:
 def make_ppo_network(
     network: networks_lib.FeedForwardNetwork, params: Dict[str, jnp.ndarray]
 ) -> PPONetworks:
-    """TODO: Add description here."""
+    """Makes generic PPO network
+
+    Args:
+        network: feedforward network representing the agent policy function
+        params: values parameterising the network.
+
+    Returns:
+        PPONetworks: PPO network class
+    """
     return PPONetworks(
         network=network,
         params=params,
@@ -119,7 +165,23 @@ def make_networks(
     critic_layer_sizes: Sequence[int] = (512, 512, 256),
     observation_network: Callable = utils.batch_concat,
 ) -> PPONetworks:
-    """TODO: Add description here."""
+    """Calls functions to make discrete or continuous network
+
+    Args:
+        spec: specifications of training environment
+        key: pseudo-random value used to initialise distributions
+        policy_layer_sizes: size of each layer of the policy network
+        critic_layer_sizes: size of each layer of the critic network
+        observation_network: Network used for feature extraction layers
+
+    Returns:
+        make_discrete_networks: function to create a discrete network
+        make_continuous_networks: function to create a continuous network
+
+    Raises:
+        NotImplementedError: Raises an error if continous network is not
+                        available
+    """
     if isinstance(spec.actions, specs.DiscreteArray):
         return make_discrete_networks(
             environment_spec=spec,
@@ -145,7 +207,18 @@ def make_discrete_networks(
     observation_network: Callable = utils.batch_concat,
     # default behaviour is to flatten observations
 ) -> PPONetworks:
-    """TODO: Add description here."""
+    """Make discrete PPO network
+
+    Args:
+        environment_spec: specifications of training environment
+        key: pseudo-random value used to initialise distributions
+        policy_layer_sizes: size of each layer of the policy network
+        critic_layer_sizes: size of each layer of the critic network
+        observation_network: Network used for feature extraction layers
+
+    Returns:
+        make_ppo_network: function to create a ppo network
+    """
 
     num_actions = environment_spec.actions.num_values
 
@@ -189,7 +262,20 @@ def make_default_networks(
     critic_layer_sizes: Sequence[int] = (512, 512, 256),
     observation_network: Callable = utils.batch_concat,
 ) -> Dict[str, Any]:
-    """Description here"""
+    """Call to create one of default Mava network types
+
+    Args:
+        environment_spec: specifications of training environment
+        agent_net_keys: keys for each agent network
+        rng_key: pseudo-random value used to initialise distributions
+        net_spec_keys: keys for each agent network
+        policy_layer_sizes: size of each layer of the policy network
+        critic_layer_sizes: size of each layer of the critic network
+        observation_network: Network used for feature extraction layers
+
+    Returns:
+        networks: networks created to given spec
+    """
 
     # Create agent_type specs.
     specs = environment_spec.get_agent_environment_specs()
diff --git a/mava/systems/jax/mappo/system.py b/mava/systems/jax/mappo/system.py
--- a/mava/systems/jax/mappo/system.py
+++ b/mava/systems/jax/mappo/system.py
@@ -26,10 +26,14 @@
 
 class MAPPOSystem(System):
     def design(self) -> Tuple[DesignSpec, Any]:
-        """Mock system design with zero components.
+        """System design for IPPO with single optimiser.
+
+        Args:
+            None.
 
         Returns:
-            system callback components
+            system: design spec for IPPO
+            default_params: default IPPO configuration
         """
         # Set the default configs
         default_params = MAPPODefaultConfig()


Test patch: