fix: remove descriptors squeeze() to fix 1D-BD tasks

qdax/core/containers/mome_repertoire.py

@@ -288,10 +288,10 @@ def _add_one(
                 cell_descriptor,
                 cell_mask,
             ) = self._update_masked_pareto_front(
-                pareto_front_fitnesses=cell_fitness.squeeze(),
-                pareto_front_genotypes=cell_genotype.squeeze(),
-                pareto_front_descriptors=cell_descriptor.squeeze(),
-                mask=cell_mask.squeeze(),
+                pareto_front_fitnesses=cell_fitness.squeeze(axis=0),
+                pareto_front_genotypes=cell_genotype.squeeze(axis=0),
+                pareto_front_descriptors=cell_descriptor.squeeze(axis=0),
+                mask=cell_mask.squeeze(axis=0),
                 new_batch_of_fitnesses=jnp.expand_dims(fitness, axis=0),
                 new_batch_of_genotypes=jnp.expand_dims(genotype, axis=0),
                 new_batch_of_descriptors=jnp.expand_dims(descriptors, axis=0),
@@ -307,6 +307,7 @@ def _add_one(
             )
             new_fitnesses = carry.fitnesses.at[index].set(cell_fitness)
             new_descriptors = carry.descriptors.at[index].set(cell_descriptor)
+
             carry = carry.replace(  # type: ignore
                 genotypes=new_genotypes,
                 descriptors=new_descriptors,

qdax/core/containers/repertoire.py

@@ -292,13 +292,13 @@ def add(
             batch_of_fitnesses.squeeze()
         )
         new_descriptors = self.descriptors.at[batch_of_indices.squeeze()].set(
-            batch_of_descriptors.squeeze()
+            batch_of_descriptors
         )
 
         return MapElitesRepertoire(
             genotypes=new_repertoire_genotypes,
             fitnesses=new_fitnesses.squeeze(),
-            descriptors=new_descriptors.squeeze(),
+            descriptors=new_descriptors,
             centroids=self.centroids,
         )
 

qdax/core/sac.py

@@ -389,7 +389,7 @@ def _update_critic(
         critic_params = optax.apply_updates(
             training_state.critic_params, critic_updates
         )
-        target_critic_params = jax.tree_multimap(
+        target_critic_params = jax.tree_map(
             lambda x1, x2: (1.0 - self._config.tau) * x1 + self._config.tau * x2,
             training_state.target_critic_params,
             critic_params,

qdax/core/td3.py

@@ -303,7 +303,7 @@ def update(
             training_state.critic_params, critic_updates
         )
         # Soft update of target critic network
-        target_critic_params = jax.tree_multimap(
+        target_critic_params = jax.tree_map(
             lambda x1, x2: (1.0 - self._config.soft_tau_update) * x1
             + self._config.soft_tau_update * x2,
             training_state.target_critic_params,
@@ -325,7 +325,7 @@ def update_policy_step() -> Tuple[Params, Params, optax.OptState]:
                 training_state.policy_params, policy_updates
             )
             # Soft update of target policy
-            target_policy_params = jax.tree_multimap(
+            target_policy_params = jax.tree_map(
                 lambda x1, x2: (1.0 - self._config.soft_tau_update) * x1
                 + self._config.soft_tau_update * x2,
                 training_state.target_policy_params,

tests/core_test/map_elites_test.py

@@ -18,9 +18,10 @@
 from qdax.types import EnvState, Params, RNGKey
 
 
-def test_map_elites() -> None:
+@pytest.mark.parametrize("environment_name", ["walker2d_uni", "hopper_uni"])
+def test_map_elites(environment_name: str) -> None:
     batch_size = 10
-    env_name = "walker2d_uni"
+    env_name = environment_name
     episode_length = 100
     num_iterations = 5
     seed = 42
@@ -152,4 +153,4 @@ def metrics_fn(repertoire: MapElitesRepertoire) -> Dict:
 
 
 if __name__ == "__main__":
-    test_map_elites()
+    test_map_elites(environment_name="walker2d_uni")

tests/core_test/mome_test.py

@@ -18,12 +18,15 @@
 from qdax.utils.metrics import compute_moqd_metrics
 
 
-def test_mome() -> None:
+@pytest.mark.parametrize("num_descriptors", [1, 2])
+def test_mome(num_descriptors: int) -> None:
 
     pareto_front_max_length = 50
     num_variables = 100
     num_iterations = 100
 
+    num_descriptors = num_descriptors
+
     num_centroids = 64
     minval = -2
     maxval = 4
@@ -41,7 +44,7 @@ def rastrigin_scorer(
         """
         Rastrigin Scorer with first two dimensions as descriptors
         """
-        descriptors = genotypes[:, :2]
+        descriptors = genotypes[:, :num_descriptors]
         f1 = -(
             10 * genotypes.shape[1]
             + jnp.sum(
@@ -79,7 +82,7 @@ def scoring_fn(
     random_key = jax.random.PRNGKey(42)
     random_key, subkey = jax.random.split(random_key)
     init_genotypes = jax.random.uniform(
-        random_key,
+        subkey,
         (batch_size, num_variables),
         minval=minval,
         maxval=maxval,
@@ -109,7 +112,7 @@ def scoring_fn(
     )
 
     centroids = compute_cvt_centroids(
-        num_descriptors=2,
+        num_descriptors=num_descriptors,
         num_init_cvt_samples=20000,
         num_centroids=num_centroids,
         minval=minval,
@@ -138,4 +141,4 @@ def scoring_fn(
 
 
 if __name__ == "__main__":
-    test_mome()
+    test_mome(num_descriptors=1)