chore: final touchs

instadeepai · Nov 11, 2024 · 369b484 · 369b484
1 parent 0adf61d
commit 369b484
Show file tree

Hide file tree

Showing 10 changed files with 11 additions and 11 deletions.
diff --git a/mava/configs/default/ff_hasac.yaml b/mava/configs/default/ff_hasac.yaml
@@ -4,7 +4,7 @@ defaults:
   - arch: anakin
   - system: sac/ff_hasac
   - network: mlp  # [mlp, cnn]
-  - env: mabrax  # [mabrax]
+  - env: mabrax  # [mabrax, mpe]
 
 hydra:
   searchpath:

diff --git a/mava/configs/default/ff_ippo.yaml b/mava/configs/default/ff_ippo.yaml
@@ -3,7 +3,7 @@ defaults:
   - arch: anakin
   - system: ppo/ff_ippo
   - network: mlp  # [mlp, cnn]
-  - env: rware  # [cleaner, connector, gigastep, lbf, mabrax, matrax, rware, smax]
+  - env: rware  # [cleaner, connector, gigastep, lbf, mabrax, matrax, rware, smax, mpe]
   - _self_
 
 hydra:

diff --git a/mava/configs/default/ff_isac.yaml b/mava/configs/default/ff_isac.yaml
@@ -4,7 +4,7 @@ defaults:
   - arch: anakin
   - system: sac/ff_isac
   - network: mlp
-  - env: mabrax  # [mabrax]
+  - env: mabrax  # [mabrax, mpe]
 
 hydra:
   searchpath:

diff --git a/mava/configs/default/ff_mappo.yaml b/mava/configs/default/ff_mappo.yaml
@@ -3,7 +3,7 @@ defaults:
   - arch: anakin
   - system: ppo/ff_mappo
   - network: mlp  # [mlp, cnn]
-  - env: rware  # [cleaner, connector, gigastep, lbf, mabrax, matrax, rware, smax]
+  - env: rware  # [cleaner, connector, gigastep, lbf, mabrax, matrax, rware, smax, mpe]
   - _self_
 
 hydra:

diff --git a/mava/configs/default/ff_masac.yaml b/mava/configs/default/ff_masac.yaml
@@ -4,7 +4,7 @@ defaults:
   - arch: anakin
   - system: sac/ff_masac
   - network: mlp
-  - env: mabrax  # [mabrax]
+  - env: mabrax  # [mabrax, mpe]
 
 hydra:
   searchpath:

diff --git a/mava/configs/default/ff_sable.yaml b/mava/configs/default/ff_sable.yaml
@@ -3,7 +3,7 @@ defaults:
   - arch: anakin
   - system: sable/ff_sable
   - network: ff_retention
-  - env: rware  # [cleaner, connector, gigastep, lbf, rware, smax]
+  - env: rware  # [cleaner, connector, gigastep, lbf, rware, smax, mpe]
   - _self_
 
 hydra:

diff --git a/mava/configs/default/mat.yaml b/mava/configs/default/mat.yaml
@@ -3,7 +3,7 @@ defaults:
   - arch: anakin
   - system: mat/mat
   - network: transformer
-  - env: rware # [gigastep, lbf, mabrax, matrax, rware, smax]
+  - env: rware # [gigastep, lbf, mabrax, matrax, rware, smax, mpe]
   - _self_
 
 hydra:

diff --git a/mava/configs/default/rec_iql.yaml b/mava/configs/default/rec_iql.yaml
@@ -4,7 +4,7 @@ defaults:
   - arch: anakin
   - system: q_learning/rec_iql
   - network: rnn  # [rnn, rcnn]
-  - env: smax  # [cleaner, connector, gigastep, lbf, matrax, rware, smax, mpe]
+  - env: smax  # [cleaner, connector, gigastep, lbf, matrax, rware, smax]
 
 hydra:
   searchpath:

diff --git a/mava/wrappers/jaxmarl.py b/mava/wrappers/jaxmarl.py
@@ -440,5 +440,5 @@ def action_mask(self, wrapped_env_state: Any) -> Array:
 
     def get_global_state(self, wrapped_env_state: Any, obs: Dict[str, Array]) -> Array:
         """Get global state from observation and copy it for each agent."""
-        all_obs = jnp.array([obs[agent] for agent in self._env.agents]).flatten()
-        return jnp.tile(all_obs, (self.num_agents, 1))
+        global_state = jnp.concatenate([obs[agent_id] for agent_id in obs])
+        return jnp.tile(global_state, (self.num_agents, 1))
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -8,7 +8,7 @@ hydra-core==1.3.2
 id-marl-eval @ git+https://github.com/instadeepai/marl-eval
 jax==0.4.30
 jaxlib==0.4.30
-jaxmarl @ git+https://github.com/RuanJohn/JaxMARL@unpin-jax
+jaxmarl @ git+https://github.com/RuanJohn/JaxMARL@unpin-jax  # This only unpins the version of Jax.
 jumanji @ git+https://github.com/sash-a/jumanji@old_jumanji  # Includes a few extra MARL envs
 matrax @ git+https://github.com/instadeepai/matrax@4c5d8aa97214848ea659274f16c48918c13e845b
 mujoco==3.1.3