feat: Upgrade Library Versions and Python Version (#187)

- moving to Python 3.10 - Upgrade all library versions in requirements.txt and setup.py - Remove DM-Haiku cause it is now deprecated. - all networks now are based on a single MLP class - jax.tree_map has been replaced with jax.tree_util.tree_map everywhere to avoid Deprecation Warnings. - types -> custom_types - add extra_require for jax[cuda12] - fix all notebooks and update typing extensions (for running notebooks) - fix dependabot security issues - added instructions for pip install qdax[cuda12] in README - fix observation space in jumanji test script
adaptive-intelligent-robotics · Aug 30, 2024 · f2c6311 · f2c6311
1 parent fe5c66a
commit f2c6311
Show file tree

Hide file tree

Showing 137 changed files with 857 additions and 733 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,17 +1,17 @@
 repos:
 -   repo: https://github.com/pycqa/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - id: isort
         args: ["--profile", "black"]
 -   repo: https://github.com/ambv/black
-    rev: 22.3.0
+    rev: 24.8.0
     hooks:
     - id: black
-      language_version: python3.9
-      args: ["--target-version", "py39"]
+      language_version: python3.10
+      args: ["--target-version", "py310"]
 -   repo: https://github.com/PyCQA/flake8
-    rev: 3.8.4
+    rev: 7.1.1
     hooks:
     - id: flake8
       args: ['--max-line-length=88', '--extend-ignore=E203']
@@ -21,12 +21,12 @@ repos:
           - flake8-comprehensions
           - flake8-bugbear
 -   repo: https://github.com/kynan/nbstripout
-    rev: 0.3.9
+    rev: 0.7.1
     hooks:
       - id: nbstripout
         args: ["examples/"]
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.0.1
+    rev: v4.6.0
     hooks:
       - id: debug-statements
       - id: requirements-txt-fixer
@@ -42,6 +42,6 @@ repos:
       - id: trailing-whitespace # This hook trims trailing whitespace
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.942
+    rev: v1.11.2
     hooks:
       - id: mypy
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
@@ -8,7 +8,7 @@ version: 2
 build:
   os: ubuntu-20.04
   tools:
-    python: "3.9"
+    python: "3.10"
   apt_packages:
     - swig
 

diff --git a/README.md b/README.md
@@ -24,6 +24,12 @@ QDax is available on PyPI and can be installed with:
 ```bash
 pip install qdax
 ```
+
+To install QDax with CUDA 12 support, use:
+```bash
+pip install qdax[cuda12]
+```
+
 Alternatively, the latest commit of QDax can be installed directly from source with:
 ```bash
 pip install git+https://github.com/adaptive-intelligent-robotics/QDax.git@main

diff --git a/dev.Dockerfile b/dev.Dockerfile
@@ -16,7 +16,7 @@ RUN micromamba create -y --file /tmp/environment.yaml \
 
 
 FROM python as test-image
-ENV PATH=/opt/conda/envs/qdaxpy39/bin/:$PATH APP_FOLDER=/app
+ENV PATH=/opt/conda/envs/qdaxpy310/bin/:$PATH APP_FOLDER=/app
 ENV PYTHONPATH=$APP_FOLDER:$PYTHONPATH
 
 COPY --from=conda /opt/conda/envs/. /opt/conda/envs/
@@ -26,7 +26,7 @@ RUN pip install -r requirements-dev.txt
 
 
 FROM nvidia/cuda:11.5.2-cudnn8-devel-ubuntu20.04 as cuda-image
-ENV PATH=/opt/conda/envs/qdaxpy39/bin/:$PATH APP_FOLDER=/app
+ENV PATH=/opt/conda/envs/qdaxpy310/bin/:$PATH APP_FOLDER=/app
 ENV PYTHONPATH=$APP_FOLDER:$PYTHONPATH
 
 
@@ -70,7 +70,7 @@ RUN apt-get update && \
     libosmesa6-dev \
     patchelf \
     python3-opengl \
-    python3-dev=3.9* \
+    python3-dev=3.10* \
     python3-pip \
     screen \
     sudo \

diff --git a/docs/installation.md b/docs/installation.md
@@ -86,7 +86,7 @@ git clone [email protected]:adaptive-intelligent-robotics/QDax.git
 
 2. Activate the environment and manually install the package qdax
     ```zsh
-    conda activate qdaxpy39
+    conda activate qdaxpy310
     pip install -e .
     ```
 

diff --git a/environment.yaml b/environment.yaml
@@ -1,9 +1,9 @@
-name: qdaxpy39
+name: qdaxpy310
 channels:
 - defaults
 - conda-forge
 dependencies:
-- python=3.9
+- python=3.10
 - pip>=20.3.3
 - conda>=4.9.2
 - pip:

diff --git a/examples/aurora.ipynb b/examples/aurora.ipynb
@@ -512,11 +512,8 @@
   }
  ],
  "metadata": {
-  "interpreter": {
-   "hash": "9ae46cf6a59eb5e192bc4f27fbb5c33d8a30eb9acb43edbb510eeaf7c819ab64"
-  },
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "venv",
    "language": "python",
    "name": "python3"
   },
@@ -530,7 +527,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,

diff --git a/examples/cmaes.ipynb b/examples/cmaes.ipynb
@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "5c4ab97a",
+   "id": "0",
    "metadata": {},
    "source": [
     "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/examples/cmaes.ipynb)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "222bbe00",
+   "id": "1",
    "metadata": {},
    "source": [
     "# Optimizing with CMA-ES in Jax\n",
@@ -26,7 +26,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d731f067",
+   "id": "2",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -71,7 +71,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "7b6e910b",
+   "id": "3",
    "metadata": {},
    "source": [
     "## Set the hyperparameters"
@@ -80,7 +80,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "404fb0dc",
+   "id": "4",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -98,7 +98,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ccc7cbeb",
+   "id": "5",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -111,7 +111,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "436dccbb",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -133,7 +133,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "62bdd2a4",
+   "id": "7",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -146,7 +146,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4cf03f55",
+   "id": "8",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -167,7 +167,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f1f69f50",
+   "id": "9",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -180,7 +180,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1a95b74d",
+   "id": "10",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "ac2d5c0d",
+   "id": "11",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -207,7 +207,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "363198ca",
+   "id": "12",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -245,7 +245,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "0e5820b8",
+   "id": "13",
    "metadata": {},
    "source": [
     "## Check final fitnesses and distribution mean"
@@ -254,7 +254,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "1e4a2c7b",
+   "id": "14",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -272,7 +272,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f3bd2b0f",
+   "id": "15",
    "metadata": {
     "pycharm": {
      "name": "#%% md\n"
@@ -285,7 +285,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ad85551c",
+   "id": "16",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
@@ -333,7 +333,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.10.12"
   },
   "vscode": {
    "interpreter": {

diff --git a/examples/cmame.ipynb b/examples/cmame.ipynb
@@ -141,9 +141,9 @@
     "def clip(x: jnp.ndarray):\n",
     "    in_bound = (x <= maxval) * (x >= minval)\n",
     "    return jnp.where(\n",
-    "        condition=in_bound,\n",
-    "        x=x,\n",
-    "        y=(maxval / x)\n",
+    "        in_bound,\n",
+    "        x,\n",
+    "        (maxval / x)\n",
     "    )\n",
     "\n",
     "def _behavior_descriptor_1(x: jnp.ndarray):\n",
@@ -387,7 +387,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.10.12"
   },
   "vscode": {
    "interpreter": {

diff --git a/examples/cmamega.ipynb b/examples/cmamega.ipynb
@@ -315,7 +315,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.13"
+   "version": "3.10.12"
   },
   "vscode": {
    "interpreter": {

diff --git a/examples/dads.ipynb b/examples/dads.ipynb
@@ -67,12 +67,6 @@
     "    import jumanji\n",
     "\n",
     "try:\n",
-    "    import haiku\n",
-    "except:\n",
-    "    !pip install git+https://github.com/deepmind/[email protected] |tail -n 1\n",
-    "    import haiku\n",
-    "\n",
-    "try:\n",
     "    import qdax\n",
     "except:\n",
     "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
@@ -554,7 +548,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.2"
+   "version": "3.10.12"
   },
   "vscode": {
    "interpreter": {

diff --git a/examples/diayn.ipynb b/examples/diayn.ipynb
@@ -67,12 +67,6 @@
     "    import jumanji\n",
     "\n",
     "try:\n",
-    "    import haiku\n",
-    "except:\n",
-    "    !pip install git+https://github.com/deepmind/[email protected] |tail -n 1\n",
-    "    import haiku\n",
-    "\n",
-    "try:\n",
     "    import qdax\n",
     "except:\n",
     "    !pip install --no-deps git+https://github.com/adaptive-intelligent-robotics/QDax@main |tail -n 1\n",
@@ -544,7 +538,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.2"
+   "version": "3.10.12"
   },
   "vscode": {
    "interpreter": {