Repository: instadeepai/Mava

Issue:
[FEATURE] MPE sum episode return
### Feature
Add additional running statistics to the environment which will be the sum_episode_return.

### Proposal
Some papers use MPE with a global return equivalent to **the sum** instead of **the mean** of all the agents reward, therefore, to make our results comparable to the existing results we need to implement this feature.

### Definition of done
Implement the sum of episode return in `mava/wrappers/environment_loop_wrappers.py`

### 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.
* [ ] Link experiment/benchmarking after implementation (optional).

<!-- Base checklist. Don’t hesitate to adapt it to your use-case. -->

### Links / references / screenshots
<!-- For ex: link to CI job, experiment runs, etc. -->



Discussion:

Resolving patch:
diff --git a/examples/petting_zoo/simple_spread/feedforward/decentralised/run_ippo.py b/examples/petting_zoo/simple_spread/feedforward/decentralised/run_ippo.py
new file mode 100644
--- /dev/null
+++ b/examples/petting_zoo/simple_spread/feedforward/decentralised/run_ippo.py
@@ -0,0 +1,116 @@
+# 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.
+
+"""Example running IPPO on MPE environments."""
+import functools
+from datetime import datetime
+from typing import Any
+
+import optax
+from absl import app, flags
+
+from mava.systems import ippo
+from mava.utils.environments import pettingzoo_utils
+from mava.utils.loggers import logger_utils
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "env_class",
+    "mpe",
+    "Pettingzoo environment class, e.g. atari (str).",
+)
+flags.DEFINE_string(
+    "env_name",
+    "simple_spread_v2",
+    "Pettingzoo environment name, e.g. pong (str).",
+)
+
+flags.DEFINE_string(
+    "mava_id",
+    str(datetime.now()),
+    "Experiment identifier that can be used to continue experiments.",
+)
+
+flags.DEFINE_string("base_dir", "~/mava", "Base dir to store experiments.")
+
+
+def main(_: Any) -> None:
+    """Run main script"""
+    # Environment.
+    environment_factory = functools.partial(
+        pettingzoo_utils.make_environment,
+        env_class=FLAGS.env_class,
+        env_name=FLAGS.env_name,
+        concat_agent_id=True,  # Concatenate agent_ids to the environment observations.
+    )
+
+    # Networks.
+    def network_factory(*args: Any, **kwargs: Any) -> Any:
+        return ippo.make_default_networks(  # type: ignore
+            policy_layer_sizes=(64,),
+            critic_layer_sizes=(256,),
+            *args,
+            **kwargs,
+        )
+
+    # Used for checkpoints, tensorboard logging and env monitoring
+    experiment_path = f"{FLAGS.base_dir}/{FLAGS.mava_id}"
+
+    # Log every [log_every] seconds.
+    log_every = 10
+    logger_factory = functools.partial(
+        logger_utils.make_logger,
+        directory=FLAGS.base_dir,
+        to_terminal=True,
+        to_tensorboard=True,
+        time_stamp=FLAGS.mava_id,
+        time_delta=log_every,
+    )
+
+    # Optimisers.
+    policy_optimiser = optax.chain(
+        optax.clip_by_global_norm(40.0), optax.scale_by_adam(), optax.scale(-1e-4)
+    )
+
+    critic_optimiser = optax.chain(
+        optax.clip_by_global_norm(40.0), optax.scale_by_adam(), optax.scale(-1e-4)
+    )
+
+    # Create the system.
+    system = ippo.IPPOSystem()
+
+    # Build the system.
+    system.build(
+        environment_factory=environment_factory,
+        network_factory=network_factory,
+        logger_factory=logger_factory,
+        experiment_path=experiment_path,
+        policy_optimiser=policy_optimiser,
+        critic_optimiser=critic_optimiser,
+        run_evaluator=True,
+        sample_batch_size=5,
+        num_epochs=15,
+        num_executors=1,
+        multi_process=True,
+        evaluation_interval={"executor_steps": 10000},
+        evaluation_duration={"evaluator_episodes": 32},
+    )
+
+    # Launch the system.
+    system.launch()
+
+
+if __name__ == "__main__":
+    app.run(main)
diff --git a/mava/wrappers/environment_loop_wrappers.py b/mava/wrappers/environment_loop_wrappers.py
--- a/mava/wrappers/environment_loop_wrappers.py
+++ b/mava/wrappers/environment_loop_wrappers.py
@@ -97,6 +97,7 @@ def __init__(
             self.__setattr__(f"_{metric}_stats", RunningStatistics(metric))
             for stat in self._summary_stats:
                 self._running_statistics[f"{stat}_{metric}"] = 0.0
+        self._running_statistics["raw_sum_episode_return"] = 0.0
 
     def _compute_step_statistics(self, rewards: Dict[str, float]) -> None:
         pass
@@ -229,7 +230,7 @@ def _compute_episode_statistics(
         # Collect the results and combine with counts.
         steps_per_second = episode_steps / (time.time() - start_time)
         mean_episode_return = np.mean(np.array(list(episode_returns.values())))
-
+        raw_sum_episode_return = np.sum(np.array(list(episode_returns.values())))
         # Record counts.
         if hasattr(self._executor, "_counts"):
             loop_type = "evaluator" if "_" not in self._loop_label else "executor"
@@ -254,6 +255,7 @@ def _compute_episode_statistics(
         self._episode_length_stats.push(episode_steps)
         self._episode_return_stats.push(mean_episode_return)
         self._steps_per_second_stats.push(steps_per_second)
+        self._running_statistics["raw_sum_episode_return"] = raw_sum_episode_return
 
         for metric in self._metrics:
             for stat in self._summary_stats:


Test patch: