Move cmaes.py from core to baselines

docs/api_documentation/core/cmaes.md

@@ -1,3 +1,3 @@
 # CMAES class
 
-::: qdax.core.cmaes.CMAES
+::: qdax.baselines.cmaes.CMAES

examples/cmaes.ipynb

@@ -55,7 +55,7 @@
     "import matplotlib.pyplot as plt\n",
     "from matplotlib.patches import Ellipse\n",
     "\n",
-    "from qdax.core.cmaes import CMAES"
+    "from qdax.baselines.cmaes import CMAES"
    ]
   },
   {

mkdocs.yml

@@ -154,11 +154,11 @@ nav:
         - DADS: api_documentation/core/dads.md
         - SAC: api_documentation/core/sac.md
         - TD3: api_documentation/core/td3.md
-        - CMAES: api_documentation/core/cmaes.md
         - Genetic Algorithm: api_documentation/core/genetic_algorithm.md
         - NSGA2: api_documentation/core/nsga2.md
         - SPEA2: api_documentation/core/spea2.md
         - PBT: api_documentation/core/pbt.md
+        - CMAES: api_documentation/core/cmaes.md
       - Containers: api_documentation/core/containers.md
       - Emitters: api_documentation/core/emitters.md
       - Neuroevolution: api_documentation/core/neuroevolution.md

qdax/core/containers/unstructured_repertoire.py

@@ -110,9 +110,7 @@ def intra_batch_comp(
 
     # We want to eliminate the same individual (distance 0)
     fitness = jnp.where(knn_relevant_indices == current_index, False, fitness)
-    current_fitness = jnp.squeeze(
-        eval_scores.at[knn_relevant_indices.at[0].get()].get()
-    )
+    current_fitness = jnp.squeeze(eval_scores.at[current_index].get())
 
     # Is the fitness of the other individual higher?
     # If both are True then we discard the current individual since this individual

qdax/core/emitters/cma_emitter.py

@@ -7,7 +7,7 @@
 import jax
 import jax.numpy as jnp
 
-from qdax.core.cmaes import CMAES, CMAESState
+from qdax.baselines.cmaes import CMAES, CMAESState
 from qdax.core.containers.mapelites_repertoire import (
     MapElitesRepertoire,
     get_cells_indices,

qdax/core/emitters/cma_mega_emitter.py

@@ -6,7 +6,7 @@
 import jax
 import jax.numpy as jnp
 
-from qdax.core.cmaes import CMAES, CMAESState
+from qdax.baselines.cmaes import CMAES, CMAESState
 from qdax.core.containers.mapelites_repertoire import (
     MapElitesRepertoire,
     get_cells_indices,

qdax/core/emitters/cma_rnd_emitter.py

@@ -6,7 +6,7 @@
 import jax
 import jax.numpy as jnp
 
-from qdax.core.cmaes import CMAESState
+from qdax.baselines.cmaes import CMAESState
 from qdax.core.containers.mapelites_repertoire import MapElitesRepertoire
 from qdax.core.emitters.cma_emitter import CMAEmitter, CMAEmitterState
 from qdax.custom_types import Descriptor, ExtraScores, Fitness, Genotype, RNGKey

tests/core_test/cmaes_test.py

@@ -4,7 +4,7 @@
 import jax.numpy as jnp
 import pytest
 
-from qdax.core.cmaes import CMAES
+from qdax.baselines.cmaes import CMAES
 
 
 def test_cmaes() -> None: