feat(notebooks): add example notebook for NSGA2 and SPEA2 (#58)

Add example notebook for NSGA2 and SPEA2 + minor fixes in MOME notebook

README.md

@@ -59,6 +59,7 @@ repertoire.genotypes, repertoire.fitnesses, repertoire.descriptors
 
 ## QDax core algorithms
 QDax currently supports the following algorithms:
+
 | Algorithm  | Example |
 | --- | --- |
 | [MAP-Elites](https://arxiv.org/abs/1504.04909) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/mapelites_example.ipynb) |
@@ -68,13 +69,17 @@ QDax currently supports the following algorithms:
 | [CMA-MEGA](https://arxiv.org/abs/2106.03894) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/cmamega_example.ipynb) |
 | [Multi-Objective Quality-Diversity (MOME)](https://arxiv.org/abs/2202.03057) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/mome_example.ipynb) |
 
+
 ## QDax baseline algorithms
 The QDax library also provides implementations for some useful baseline algorithms:
+
 | Algorithm  | Example |
 | --- | --- |
 | [DIAYN](https://arxiv.org/abs/1802.06070) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/diayn_example.ipynb) |
 | [DADS](https://arxiv.org/abs/1907.01657) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/dads_example.ipynb) |
 | [SMERL](https://arxiv.org/abs/2010.14484) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/smerl_example.ipynb) |
+| [NSGA2](https://ieeexplore.ieee.org/document/996017) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/nsga2_spea2_example.ipynb) |
+| [SPEA2](https://www.semanticscholar.org/paper/SPEA2%3A-Improving-the-strength-pareto-evolutionary-Zitzler-Laumanns/b13724cb54ae4171916f3f969d304b9e9752a57f) | [![Open All Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/adaptive-intelligent-robotics/QDax/blob/main/nsga2_spea2_example.ipynb) |
 
 
 ## QDax Overview

notebooks/mome_example.ipynb

@@ -86,20 +86,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "pareto_front_max_length = 50\n",
-    "num_variables = 100\n",
-    "num_iterations = 1000\n",
+    "#@markdown ---\n",
+    "pareto_front_max_length = 50 #@param {type:\"integer\"}\n",
+    "num_variables = 100 #@param {type:\"integer\"}\n",
+    "num_iterations = 1000 #@param {type:\"integer\"}\n",
     "\n",
-    "num_centroids = 64\n",
-    "minval = -2\n",
-    "maxval = 4\n",
-    "proportion_to_mutate = 0.6\n",
-    "eta = 1\n",
-    "proportion_var_to_change = 0.5\n",
-    "crossover_percentage = 1.\n",
-    "batch_size = 100\n",
-    "lag = 2.2\n",
-    "base_lag = 0"
+    "num_centroids = 64 #@param {type:\"integer\"}\n",
+    "minval = -2 #@param {type:\"number\"}\n",
+    "maxval = 4 #@param {type:\"number\"}\n",
+    "proportion_to_mutate = 0.6 #@param {type:\"number\"}\n",
+    "eta = 1 #@param {type:\"number\"}\n",
+    "proportion_var_to_change = 0.5 #@param {type:\"number\"}\n",
+    "crossover_percentage = 1. #@param {type:\"number\"}\n",
+    "batch_size = 100 #@param {type:\"integer\"}\n",
+    "lag = 2.2 #@param {type:\"number\"}\n",
+    "base_lag = 0 #@param {type:\"number\"}\n",
+    "#@markdown ---"
    ]
   },
   {
@@ -120,8 +122,8 @@
    "outputs": [],
    "source": [
     "def rastrigin_scorer(\n",
-    "    genotypes: jnp.ndarray, base_lag: int, lag: int\n",
-    ") -> Tuple[jnp.ndarray, jnp.ndarray]:\n",
+    "    genotypes: jnp.ndarray, base_lag: float, lag: float\n",
+    ") -> Tuple[Fitness, Descriptor]:\n",
     "    \"\"\"\n",
     "    Rastrigin Scorer with first two dimensions as descriptors\n",
     "    \"\"\"\n",