Merge branch 'develop' into feat/integ-tests

instadeepai · Jul 23, 2024 · 92855b2 · 92855b2
2 parents 6ef1f5e + 7267aa3
commit 92855b2
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 5 deletions.
diff --git a/.github/workflows/tests_linters.yaml b/.github/workflows/tests_linters.yaml
@@ -6,7 +6,7 @@ jobs:
   tests-and-linters:
     name: "Python ${{ matrix.python-version }} on ${{ matrix.os }}"
     runs-on: "${{ matrix.os }}"
-    timeout-minutes: 20
+    timeout-minutes: 10
 
     strategy:
       matrix:

diff --git a/Dockerfile b/Dockerfile
@@ -36,7 +36,7 @@ RUN pip install --quiet --upgrade pip setuptools wheel &&  \
 # Need to use specific cuda versions for jax
 ARG USE_CUDA=true
 RUN if [ "$USE_CUDA" = true ] ; \
-    then pip install "jax[cuda12]==0.4.26" ; \
+    then pip install "jax[cuda12]==0.4.30" ; \
     fi
 
 # Copy all code

diff --git a/docs/DETAILED_INSTALL.md b/docs/DETAILED_INSTALL.md
@@ -22,7 +22,7 @@ pip install -e .
 
 4. Install jax on your accelerator. The example below is for an NVIDIA GPU, please the [official install guide](https://github.com/google/jax#installation) for other accelerators
 ```bash
-pip install "jax[cuda12]==0.4.26"
+pip install "jax[cuda12]==0.4.30"
 ```
 
 5. Run a system!

diff --git a/mava/wrappers/jumanji.py b/mava/wrappers/jumanji.py
@@ -143,6 +143,17 @@ def modify_timestep(self, timestep: TimeStep) -> TimeStep[Observation]:
         discount = jnp.repeat(timestep.discount, self.num_agents)
         return timestep.replace(observation=observation, reward=reward, discount=discount)
 
+    def observation_spec(
+        self,
+    ) -> specs.Spec[Union[Observation, ObservationGlobalState]]:
+        # need to cast the agents view and global state to floats as we do in modify timestep
+        inner_spec = super().observation_spec()
+        spec = inner_spec.replace(agents_view=inner_spec.agents_view.replace(dtype=float))
+        if self.add_global_state:
+            spec = inner_spec.replace(global_state=inner_spec.global_state.replace(dtype=float))
+
+        return spec
+
 
 class LbfWrapper(JumanjiMarlWrapper):
     """Multi-agent wrapper for the Level-Based Foraging environment.
@@ -192,6 +203,17 @@ def modify_timestep(self, timestep: TimeStep) -> TimeStep[Observation]:
         # Aggregate the list of individual rewards and use a single team_reward.
         return self.aggregate_rewards(timestep, modified_observation)
 
+    def observation_spec(
+        self,
+    ) -> specs.Spec[Union[Observation, ObservationGlobalState]]:
+        # need to cast the agents view and global state to floats as we do in modify timestep
+        inner_spec = super().observation_spec()
+        spec = inner_spec.replace(agents_view=inner_spec.agents_view.replace(dtype=float))
+        if self.add_global_state:
+            spec = inner_spec.replace(global_state=inner_spec.global_state.replace(dtype=float))
+
+        return spec
+
 
 class ConnectorWrapper(JumanjiMarlWrapper):
     """Multi-agent wrapper for the MA Connector environment.

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -6,8 +6,8 @@ flax
 gigastep @ git+https://github.com/mlech26l/gigastep
 hydra-core==1.3.2
 id-marl-eval @ git+https://github.com/instadeepai/marl-eval
-jax==0.4.26   # Fixed due to Mabrax dependency
-jaxlib==0.4.26
+jax==0.4.30
+jaxlib==0.4.30
 jaxmarl
 jumanji @ git+https://github.com/sash-a/jumanji  # Includes a few extra MARL envs
 matrax @ git+https://github.com/instadeepai/matrax