Merge branch 'tmp' into chore/update-libs

morgangiraud · Apr 3, 2024 · a15fc51 · a15fc51
2 parents f50b3a5 + 464747b
commit a15fc51
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 11 deletions.
diff --git a/examples/qd_cmame.py b/examples/qd_cmame.py
@@ -62,7 +62,15 @@ def run(omegaConf: DictConfig) -> None:
     grid_shape = config["grid"]["shape"]
     assert len(grid_shape) == len(features_domain)
     if True:
-        archive = GridArchive(solution_dim=len(grid_shape), dims=grid_shape, ranges=features_domain, seed=seed)
+        archive = GridArchive(
+            solution_dim=len(grid_shape),
+            dims=grid_shape,
+            ranges=features_domain,
+            seed=seed,
+            extra_fields={
+                "metadata": ((), object),
+            },
+        )
     else:
         bins = math.prod(grid_shape)
         archive = CVTArchive(bins, features_domain, seed=seed, use_kd_tree=True)

diff --git a/leniax/qd.py b/leniax/qd.py
@@ -271,8 +271,7 @@ def run_qd_search(
             fits.append(ind.fitness)
             bcs.append(ind.features)
             metadata.append(ind.get_config())
-        # scheduler.tell(fits, bcs, **metadata) # TODO: check the impact of removing metadata
-        scheduler.tell(fits, bcs)
+        scheduler.tell(fits, bcs, metadata=metadata)
 
         # Log statistics
         if itr % log_freq == 0 or itr == nb_iter:
@@ -355,10 +354,10 @@ def render_best(grid: ArchiveBase, fitness_threshold: float):
     rng_key = leniax_utils.seed_everything(seed)
 
     real_bests = []
-    for idx in grid._occupied_indices:
-        if abs(grid._objective_values[idx]) >= fitness_threshold:
+    for elite in grid:
+        if abs(elite["objective"]) >= fitness_threshold:
             rng_key, subkey = jax.random.split(rng_key)
-            lenia = LeniaIndividual(grid._metadata[idx], subkey, grid._solutions[idx])
+            lenia = LeniaIndividual(elite["metadata"], subkey, elite["solution"])
             real_bests.append(lenia)
 
     logging.info(f"Found {len(real_bests)} beast!")

diff --git a/scripts/qd_rastrigin.py b/scripts/qd_rastrigin.py
@@ -44,7 +44,15 @@ def run(omegaConf: DictConfig) -> None:
     assert len(grid_shape) == len(features_domain)
 
     bins = math.prod(grid_shape)
-    archive = CVTArchive(bins, features_domain, seed=seed, use_kd_tree=True)
+    archive = CVTArchive(
+        bins,
+        features_domain,
+        seed=seed,
+        use_kd_tree=True,
+        extra_fields={
+            "metadata": ((), object),
+        },
+    )
     archive.qd_config = config
 
     # Emitters

diff --git a/tools/vis_behaviours_means.py b/tools/vis_behaviours_means.py
@@ -56,10 +56,10 @@ def run() -> None:
 
         fitness_threshold = qd_config["run_params"]["max_run_iter"]
         real_bests = []
-        for idx in grid._occupied_indices:
-            if abs(grid._objective_values[idx]) >= fitness_threshold:
+        for elite in grid:
+            if abs(elite["objective"]) >= fitness_threshold:
                 rng_key, subkey = jax.random.split(rng_key)
-                lenia = LeniaIndividual(grid._metadata[idx], subkey, grid._solutions[idx])
+                lenia = LeniaIndividual(elite["metadata"], subkey, elite["solution"])
                 real_bests.append(lenia)
 
         print(f"Found {len(real_bests)} beast in {file_path}")
@@ -89,7 +89,7 @@ def run() -> None:
             ]
             behaviour_archive.add(lenia, 1024, behaviour, config)
 
-        print(len(behaviour_archive._occupied_indices))
+        print(len(behaviour_archive))
         leniax_qd.save_heatmap(behaviour_archive, fitness_domain, f"{subdir}/behaviour_archive_heatmap.png")
         with open(f"{subdir}/behaviour_final.p", "wb") as handle:
             pickle.dump(behaviour_archive, handle, protocol=pickle.HIGHEST_PROTOCOL)