fix: fix batch size parameters in qdax (#231)

Author: hannah-jan

examples/nsga2_spea2.ipynb

@@ -98,7 +98,6 @@
     "proportion_mutation = 0.80 #@param {type:\"number\"}\n",
     "minval = -5.12 #@param {type:\"number\"}\n",
     "maxval = 5.12 #@param {type:\"number\"}\n",
-    "batch_size = 100 #@param {type:\"integer\"}\n",
     "genotype_dim = 6 #@param {type:\"integer\"}\n",
     "lag = 2.2 #@param {type:\"number\"}\n",
     "base_lag = 0 #@param {type:\"number\"}\n",
@@ -184,7 +183,7 @@
     "key = jax.random.key(0)\n",
     "key, subkey = jax.random.split(key)\n",
     "genotypes = jax.random.uniform(\n",
-    "    subkey, (batch_size, genotype_dim), minval=minval, maxval=maxval, dtype=jnp.float32\n",
+    "    subkey, (population_size, genotype_dim), minval=minval, maxval=maxval, dtype=jnp.float32\n",
     ")\n",
     "\n",
     "# Mutation & Crossover\n",
@@ -202,11 +201,12 @@
     ")\n",
     "\n",
     "# Define the emitter\n",
+    "# NSGA-II and SPEA2 use batch size = population size\n",
     "mixing_emitter = MixingEmitter(\n",
     "    mutation_fn=mutation_function, \n",
     "    variation_fn=crossover_function, \n",
     "    variation_percentage=1-proportion_mutation, \n",
-    "    batch_size=batch_size\n",
+    "    batch_size=population_size, \n",
     ")"
    ]
   },