instadeepai · KaleabTessera · Jun 15, 2022 · Sep 27, 2021 · Sep 27, 2021 · Sep 27, 2021
@@ -18,11 +18,11 @@ jobs:
     steps:
     - name: Checkout mava
       uses: actions/checkout@v2
-    - name: Run tests in docker
-      run: |
-        docker run --mount "type=bind,src=$(pwd),dst=/tmp/mava" \
-            -w "/tmp/mava" --rm ${{ matrix.docker-image }} /bin/bash bash_scripts/tests.sh
     - name: Check format and types
       run: |
         docker run --mount "type=bind,src=$(pwd),dst=/tmp/mava" \
             -w "/tmp/mava" --rm ${{ matrix.docker-image }} /bin/bash bash_scripts/check_format.sh
+    - name: Run tests in docker
+      run: |
+        docker run --mount "type=bind,src=$(pwd),dst=/tmp/mava" \
+            -w "/tmp/mava" --rm ${{ matrix.docker-image }} /bin/bash bash_scripts/tests.sh
@@ -98,6 +98,71 @@ jobs:
           tags: ${{ steps.vars2.outputs.docker_repo }}:meltingpot-latest
           target: meltingpot
 
+
+      -
+        name: Build and push jax-core
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:jax-core-latest
+          target: jax-core
+
+      -
+        name: Build and push pz-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:pz-jax-latest
+          target: pz-jax
+
+      -
+        name: Build and push sc2-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:sc2-jax-latest
+          target: sc2-jax
+
+
+      -
+        name: Build and push flatland-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:flatland-jax-latest
+          target: flatland-jax
+
+      -
+        name: Build and push robocup-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:robocup-jax-latest
+          target: robocup-jax
+
+      -
+        name: Build and push openspiel-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:openspiel-jax-latest
+          target: openspiel-jax
+
+      -
+        name: Build and push meltingpot-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:meltingpot-jax-latest
+          target: meltingpot-jax
+
   # Action pushes new docker image if we have a PR to dev and have the label benchmark
   build-push-branches:
     if:  ${{ github.event.label.name == 'benchmark' }}
@@ -196,3 +261,66 @@ jobs:
           push: true
           tags: ${{ steps.vars2.outputs.docker_repo }}:meltingpot-${{ steps.vars.outputs.sha_short }}
           target: meltingpot
+
+      -
+        name: Build and push jax-core
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:jax-core-${{ steps.vars.outputs.sha_short }}
+          target: jax-core
+
+      -
+        name: Build and push pz-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:pz-jax-${{ steps.vars.outputs.sha_short }}
+          target: pz-jax
+
+      -
+        name: Build and push sc2-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:sc2-jax-${{ steps.vars.outputs.sha_short }}
+          target: sc2-jax
+
+      -
+        name: Build and push flatland-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:flatland-jax-${{ steps.vars.outputs.sha_short }}
+          target: flatland-jax
+
+      -
+        name: Build and push robocup-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:robocup-jax-${{ steps.vars.outputs.sha_short }}
+          target: robocup-jax
+
+      -
+        name: Build and push openspiel-jax
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:openspiel-jax-${{ steps.vars.outputs.sha_short }}
+          target: openspiel-jax
+
+      -
+        name: Build and push meltingpot
+        uses: docker/build-push-action@v2
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.vars2.outputs.docker_repo }}:meltingpot-jax-${{ steps.vars.outputs.sha_short }}
+          target: meltingpot-jax
@@ -26,7 +26,7 @@ repos:
       - id: check-yaml
 
   - repo: https://github.com/psf/black
-    rev: 21.4b1
+    rev: 22.3.0
     hooks:
       - id: black
         exclude_types: [image]

@@ -1,6 +1,6 @@
 ##########################################################
 # Core Mava image
-FROM nvidia/cuda:11.4.2-cudnn8-runtime-ubuntu20.04 as mava-core
+FROM nvidia/cuda:11.5.1-cudnn8-devel-ubuntu20.04 as mava-core
 # Flag to record agents
 ARG record
 # Ensure no installs try launch interactive screen
@@ -23,13 +23,13 @@ WORKDIR ${folder}
 COPY . /home/app/mava
 # For box2d
 RUN apt-get install swig -y
-## Install core dependencies.
-RUN pip install -e .[reverb,launchpad]
+## Install core dependencies + reverb.
+RUN pip install -e .[reverb]
 ## Optional install for screen recording.
 ENV DISPLAY=:0
 RUN if [ "$record" = "true" ]; then \
     ./bash_scripts/install_record.sh; \
-fi
+    fi
 EXPOSE 6006
 ##########################################################
 
@@ -95,3 +95,67 @@ RUN ./bash_scripts/install_meltingpot.sh
 # Add meltingpot to python path
 ENV PYTHONPATH "${PYTHONPATH}:${folder}/../packages/meltingpot"
 ##########################################################
+
+# New Jax Images
+##########################################################
+# Core Mava-Jax image
+FROM mava-core as jax-core
+# Jax gpu config.
+ENV XLA_PYTHON_CLIENT_PREALLOCATE=false
+## Install core jax dependencies.
+# Install jax gpu
+RUN pip install -e .[jax]
+RUN pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
+##########################################################
+
+##########################################################
+# PZ image
+FROM jax-core AS pz-jax
+RUN pip install -e .[pz]
+# PettingZoo Atari envs
+RUN apt-get update
+RUN apt-get install ffmpeg libsm6 libxext6  -y
+RUN apt-get install -y unrar-free
+RUN pip install autorom
+RUN AutoROM -v
+##########################################################
+
+##########################################################
+# SMAC image
+FROM jax-core AS sc2-jax
+## Install smac environment
+RUN apt-get -y install git
+RUN pip install .[sc2]
+# We use the pz wrapper for smac
+RUN pip install .[pz]
+ENV SC2PATH /home/app/mava/3rdparty/StarCraftII
+##########################################################
+
+##########################################################
+# Flatland Image
+FROM jax-core AS flatland-jax
+RUN pip install -e .[flatland]
+##########################################################
+
+#########################################################
+## Robocup Image
+FROM jax-core AS robocup-jax
+RUN apt-get install sudo -y
+RUN ./bash_scripts/install_robocup.sh
+##########################################################
+
+##########################################################
+## OpenSpiel Image
+FROM jax-core AS openspiel-jax
+RUN pip install .[open_spiel]
+##########################################################
+
+##########################################################
+# MeltingPot Image
+FROM jax-core AS meltingpot-jax
+# Install meltingpot
+RUN apt-get install -y git
+RUN ./bash_scripts/install_meltingpot.sh
+# Add meltingpot to python path
+ENV PYTHONPATH "${PYTHONPATH}:${folder}/../packages/meltingpot"
+##########################################################
@@ -26,20 +26,38 @@ DOCKER_RUN=docker run $(RUN_FLAGS) $(IMAGE)
 DOCKER_RUN_TENSORBOARD=docker run $(RUN_FLAGS_TENSORBOARD) $(IMAGE)
 
 # Choose correct image for example
-ifneq (,$(findstring debugging,$(example)))
-DOCKER_IMAGE_TAG=tf-core
-else ifneq (,$(findstring petting,$(example)))
-DOCKER_IMAGE_TAG=pz
-else ifneq (,$(findstring flatland,$(example)))
-DOCKER_IMAGE_TAG=flatland
-else ifneq (,$(findstring openspiel,$(example)))
-DOCKER_IMAGE_TAG=openspiel
-else ifneq (,$(findstring robocup,$(example)))
-DOCKER_IMAGE_TAG=robocup
-else ifneq (,$(findstring smac,$(example)))
-DOCKER_IMAGE_TAG=sc2
-else ifneq (,$(findstring meltingpot,$(example)))
-DOCKER_IMAGE_TAG=meltingpot
+ifneq (,$(findstring jax,$(example)))
+	ifneq (,$(findstring debugging,$(example)))
+	DOCKER_IMAGE_TAG=jax-core
+	else ifneq (,$(findstring petting,$(example)))
+	DOCKER_IMAGE_TAG=pz-jax
+	else ifneq (,$(findstring flatland,$(example)))
+	DOCKER_IMAGE_TAG=flatland-jax
+	else ifneq (,$(findstring openspiel,$(example)))
+	DOCKER_IMAGE_TAG=openspiel-jax
+	else ifneq (,$(findstring robocup,$(example)))
+	DOCKER_IMAGE_TAG=robocup-jax
+	else ifneq (,$(findstring smac,$(example)))
+	DOCKER_IMAGE_TAG=sc2-jax
+	else ifneq (,$(findstring meltingpot,$(example)))
+	DOCKER_IMAGE_TAG=meltingpot-jax
+	endif
+else
+	ifneq (,$(findstring debugging,$(example)))
+	DOCKER_IMAGE_TAG=tf-core
+	else ifneq (,$(findstring petting,$(example)))
+	DOCKER_IMAGE_TAG=pz
+	else ifneq (,$(findstring flatland,$(example)))
+	DOCKER_IMAGE_TAG=flatland
+	else ifneq (,$(findstring openspiel,$(example)))
+	DOCKER_IMAGE_TAG=openspiel
+	else ifneq (,$(findstring robocup,$(example)))
+	DOCKER_IMAGE_TAG=robocup
+	else ifneq (,$(findstring smac,$(example)))
+	DOCKER_IMAGE_TAG=sc2
+	else ifneq (,$(findstring meltingpot,$(example)))
+	DOCKER_IMAGE_TAG=meltingpot
+	endif
 endif
 
 # make file commands

@@ -150,11 +150,20 @@ docker run --gpus all -it --rm  -v $(pwd):/home/app/mava -w /home/app/mava insta
 
     1.1 Only Mava core:
 
+    Tensorflow version:
     ```bash
-    make build
+    make build version=tf-core
+    ```
+
+    Jax version:
+    ```bash
+    make build version=jax-core
     ```
 
     1.2 For **optional** environments:
+
+    **Note for jax images, append `-jax` to the build command, e.g. `make build version=pz-jax`.**
+
     * PettingZoo:
 
         ```
@@ -254,15 +263,17 @@ docker run --gpus all -it --rm  -v $(pwd):/home/app/mava -w /home/app/mava insta
     * Install core dependencies:
 
     ```bash
-    pip install id-mava[tf,reverb,launchpad]
+    pip install id-mava[tf,reverb]
     ```
 
     * Or for the latest version of mava from source (**you can do this for all pip install commands below for the latest depedencies**):
 
     ```bash
-    pip install git+https://github.com/instadeepai/Mava#egg=id-mava[reverb,tf,launchpad]
+    pip install git+https://github.com/instadeepai/Mava#egg=id-mava[reverb,tf]
     ```
 
+    **For the jax version of mava, please replace `tf` with `jax`, e.g. `pip install id-mava[jax,reverb]`**
+
     1.2 For **optional** environments:
     * PettingZoo:
 

@@ -20,4 +20,7 @@ pip install https://github.com/deepmind/lab2d/releases/download/release_candidat
 mkdir ../packages
 cd ../packages
 git clone -b main https://github.com/deepmind/meltingpot
-pip install meltingpot/
+# Temp - Pin version before python 3.9 requirement.
+cd meltingpot/
+git checkout c4fea35d3b57b0c4fa10db823a81e16244826e8d
+pip install .
@@ -46,7 +46,7 @@ apt-get -y install git
 apt-get install swig -y
 
 # Install depedencies
-pip install .[tf,envs,reverb,testing_formatting,launchpad,record_episode]
+pip install .[tf,envs,reverb,testing_formatting,record_episode]
 
 # For atari envs
 apt-get -y install unrar-free

@@ -1,4 +1,4 @@
-git+https://github.com/instadeepai/Mava#egg=id-mava[tf,envs,reverb,launchpad]
+git+https://github.com/instadeepai/Mava#egg=id-mava[tf,envs,reverb]
 mkdocstrings
 mkdocs-material
 pymdown-extensions