fix: dependencies in notebook examples (#149)

fix: dependencies in notebook examples
adaptive-intelligent-robotics · May 17, 2023 · 1e016f1 · 1e016f1
1 parent 773e646
commit 1e016f1
Show file tree

Hide file tree

Showing 17 changed files with 117 additions and 125 deletions.
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
@@ -18,7 +18,7 @@ mkdocs:
 # Optionally declare the Python requirements required to build your docs
 python:
   install:
-    - requirements: requirements.txt
-    - requirements: docs/requirements.txt
     - method: pip
       path: .
+    - requirements: requirements.txt
+    - requirements: docs/requirements.txt
diff --git a/environment.yaml b/environment.yaml
@@ -8,6 +8,5 @@ dependencies:
 - conda>=4.9.2
 - pip:
     - --find-links https://storage.googleapis.com/jax-releases/jax_releases.html
-    - jaxlib==0.3.15
     - -r requirements.txt
     - -r requirements-dev.txt
diff --git a/examples/cmame.ipynb b/examples/cmame.ipynb
@@ -49,7 +49,13 @@
     "except:\n",
     "    !pip install --no-deps git+https://github.com/deepmind/[email protected] |tail -n 1\n",
     "    import chex\n",
-    "    \n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
     "try:\n",
     "    import qdax\n",
     "except:\n",

diff --git a/examples/cmamega.ipynb b/examples/cmamega.ipynb
@@ -43,7 +43,13 @@
     "except:\n",
     "    !pip install --no-deps git+https://github.com/deepmind/[email protected] |tail -n 1\n",
     "    import chex\n",
-    "    \n",
+    "\n",
+    "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
     "try:\n",
     "    import qdax\n",
     "except:\n",

diff --git a/examples/dads.ipynb b/examples/dads.ipynb
@@ -45,10 +45,16 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.0.15 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
     "    import haiku\n",
     "except:\n",
     "    !pip install git+https://github.com/deepmind/[email protected] |tail -n 1\n",

diff --git a/examples/diayn.ipynb b/examples/diayn.ipynb
@@ -45,10 +45,16 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.0.15 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\n",
+    "\n",
+    "try:\n",
     "    import haiku\n",
     "except:\n",
     "    !pip install git+https://github.com/deepmind/[email protected] |tail -n 1\n",

diff --git a/examples/distributed_mapelites.ipynb b/examples/distributed_mapelites.ipynb
@@ -2,22 +2,14 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "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/distributed_mapelites.ipynb)"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "# Optimizing with MAP-Elites in Jax (multi-devices example)\n",
     "\n",
@@ -34,11 +26,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "#@title Installs and Imports\n",
@@ -61,10 +49,16 @@
     "try:\n",
     "    import brax\n",
     "except:\n",
-    "    !pip install git+https://github.com/google/brax.git@v0.0.15 |tail -n 1\n",
+    "    !pip install git+https://github.com/google/brax.git@v0.1.2 |tail -n 1\n",
     "    import brax\n",
     "\n",
     "try:\n",
+    "    import jumanji\n",
+    "except:\n",
+    "    !pip install \"jumanji==0.2.2\"\n",
+    "    import jumanji\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",
@@ -93,34 +87,22 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Setup and get devices"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "Setup the default platform where the MAP-Elites will be stored and MAP-Elite updates will happen. "
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "default_device = 'cpu'\n",
@@ -130,11 +112,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Get devices (change gpu by tpu if needed)\n",
@@ -146,23 +124,15 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Setup run parameters"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "#@title QD Training Definitions Fields\n",
@@ -185,11 +155,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Init environment, policy, population params, init states of the env\n",
     "\n",
@@ -199,11 +165,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "%%time\n",
@@ -237,11 +199,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Define the way the policy interacts with the env\n",
     "\n",
@@ -251,11 +209,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Define the fonction to play a step with the policy in the environment\n",
@@ -289,11 +243,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Define the scoring function and the way metrics are computed\n",
     "\n",
@@ -303,11 +253,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Prepare the scoring function\n",
@@ -332,11 +278,7 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Define the emitter\n",
     "\n",
@@ -346,11 +288,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Define emitter\n",
@@ -367,23 +305,15 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Instantiate and initialise the MAP Elites algorithm"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "%%time\n",
@@ -423,23 +353,15 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
+   "metadata": {},
    "source": [
     "## Launch MAP-Elites iterations"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "log_period = 10\n",
@@ -493,11 +415,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Get the repertoire from the first device\n",