Repository: instadeepai/Mava

Issue:
[FEATURE] Death masking
Add death masking to the JAX IPPO system.



Discussion:

Resolving patch:
diff --git a/examples/smac/feedforward/decentralised/run_ippo_smac_death_masking.py b/examples/smac/feedforward/decentralised/run_ippo_smac_death_masking.py
new file mode 100644
--- /dev/null
+++ b/examples/smac/feedforward/decentralised/run_ippo_smac_death_masking.py
@@ -0,0 +1,106 @@
+# 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.
+"""Run IPPO on SMAC with death masking."""
+
+
+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.smac_utils import make_environment
+from mava.utils.loggers import logger_utils
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string(
+    "map_name",
+    "3m",
+    "Starcraft 2 micromanagement map name (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 IPPO on SMAC with death masking."""
+    # Environment
+    environment_factory = functools.partial(
+        make_environment, map_name=FLAGS.map_name, death_masking=True
+    )
+
+    # Networks.
+    def network_factory(*args: Any, **kwargs: Any) -> Any:
+        return ippo.make_default_networks(  # type: ignore
+            policy_layer_sizes=(64, 64),
+            critic_layer_sizes=(64, 64, 64),
+            *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,
+    )
+
+    # Launch the system.
+    system.launch()
+
+
+if __name__ == "__main__":
+    app.run(main)
diff --git a/mava/utils/environments/smac_utils.py b/mava/utils/environments/smac_utils.py
--- a/mava/utils/environments/smac_utils.py
+++ b/mava/utils/environments/smac_utils.py
@@ -39,13 +39,23 @@ def make_environment(
         concat_agent_id: bool = False,
         evaluation: bool = False,
         random_seed: Optional[int] = None,
+        death_masking: bool = False,
     ) -> Any:
+        """Create smac environment wrapper.
 
+        Args:
+            map_name: the name of the scenario
+            concat_prev_actions: Concat one-hot vector of agent prev_action to obs.
+            concat_agent_id: Concat one-hot vector of agent ID to obs.
+            evaluation: extra param for evaluation
+            random_seed: seed
+            death_masking: whether to mask out agent observations once dead
+        """
         # Env uses int64 action space due to the use of spac.Discrete.
         set_jax_double_precision()
         env = StarCraft2Env(map_name=map_name, seed=random_seed)
 
-        env = SMACWrapper(env)
+        env = SMACWrapper(env, death_masking=death_masking)
 
         if concat_prev_actions:
             env = ConcatPrevActionToObservation(env)
diff --git a/mava/wrappers/smac.py b/mava/wrappers/smac.py
--- a/mava/wrappers/smac.py
+++ b/mava/wrappers/smac.py
@@ -33,6 +33,7 @@ def __init__(
         self,
         environment: StarCraft2Env,
         return_state_info: bool = True,
+        death_masking: bool = False,
     ):
         """Constructor for parallel PZ wrapper.
 
@@ -41,6 +42,7 @@ def __init__(
             env_preprocess_wrappers (Optional[List], optional): Wrappers
                 that preprocess envs.
                 Format (env_preprocessor, dict_with_preprocessor_params).
+            death_masking: whether to mask out agent observations once dead.
             return_state_info: return extra state info
         """
         self._environment = environment
@@ -53,6 +55,8 @@ def __init__(
         self._battles_won = 0
         self._battles_game = 0
 
+        self._death_masking = death_masking
+
     def reset(self) -> dm_env.TimeStep:
         """Resets the env.
 
@@ -183,6 +187,17 @@ def _get_legal_actions(self) -> List:
             )
         return legal_actions
 
+    def is_dead(self, agent: Any) -> bool:
+        """Check if the agent is dead.
+
+        Returns:
+            is_dead: boolean indicating whether the agent is alive or dead.
+        """
+        is_dead = False
+        if self._environment.agents[agent].health == 0.0:
+            is_dead = True
+        return is_dead
+
     def _convert_observations(
         self, observations: List, legal_actions: List, done: bool
     ) -> types.Observation:
@@ -197,9 +212,14 @@ def _convert_observations(
         """
         olt_observations = {}
         for i, agent in enumerate(self._agents):
+            # Check if agent is dead, if so, apply death mask.
+            if self._death_masking and self.is_dead(i):
+                observation = np.zeros_like(observations[i])
+            else:
+                observation = observations[i]
 
             olt_observations[agent] = types.OLT(
-                observation=observations[i],
+                observation=observation,
                 legal_actions=legal_actions[i],
                 terminal=np.asarray([done], dtype=np.float32),
             )


Test patch: